Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
01ed65ac73 | ||
|
|
5fa16d4aa3 | ||
|
|
cf0090d3f5 | ||
|
|
035e816db6 | ||
|
|
f98f6fa08d | ||
|
|
ce82dfb912 | ||
|
|
5ded3189c1 | ||
|
|
f1b1d4731b | ||
|
|
e9fe586442 | ||
|
|
6d9f3fc8fb | ||
|
|
911b801883 | ||
|
|
a576cff69f | ||
|
|
cf0cef3f25 | ||
|
|
7b6430e20a | ||
|
|
611c19d2f4 | ||
|
|
2763e9f9fa | ||
|
|
3bb21c621c | ||
|
|
374ce6e7a8 | ||
|
|
e2ec6df3d4 | ||
|
|
38daedf131 | ||
|
|
1e1b92a1ba | ||
|
|
ef66d7bd3e | ||
|
|
f00d6408d8 | ||
|
|
6962fecf86 | ||
|
|
90d706c943 | ||
|
|
d17be56e1b | ||
|
|
94366938d1 | ||
|
|
e13ea3c92f | ||
|
|
da322f8c94 | ||
|
|
6c7e4d5b35 | ||
|
|
aa768f57dd | ||
|
|
f45da0e1ce | ||
|
|
d70dab56e5 | ||
|
|
772f1639a9 | ||
|
|
c01d8e3c12 | ||
|
|
96f19636e6 | ||
|
|
8146250274 | ||
|
|
00284d92da | ||
|
|
6fc6e683eb | ||
|
|
2cf9f42e73 | ||
|
|
7f45bc6df7 | ||
|
|
ad32a51497 | ||
|
|
0e95dc3431 | ||
|
|
8dfa290417 | ||
|
|
ff090c5ebb | ||
|
|
d0c747a609 | ||
|
|
37962e30eb | ||
|
|
a22272f7a5 | ||
|
|
fdda1c4ad5 | ||
|
|
42f9c344f8 | ||
|
|
6ea8c591b7 | ||
|
|
d60e4ca52f | ||
|
|
30f83f2401 | ||
|
|
6a5834e243 | ||
|
|
0a2c274a27 | ||
|
|
07f6087963 | ||
|
|
f7e6fed1af | ||
|
|
9e8864e453 | ||
|
|
0e527fd80a | ||
|
|
0cb941e894 | ||
|
|
e1f1121cb8 | ||
|
|
efe2138e4d | ||
|
|
b616a26df1 | ||
|
|
8652bb70d6 | ||
|
|
d0aa7d7080 | ||
|
|
15bbfc4044 | ||
|
|
75b0ebe69c | ||
|
|
c9ca4be8b8 | ||
|
|
7f92c4d5b8 | ||
|
|
28c386d529 | ||
|
|
bb34b9cf3b | ||
|
|
be205f3dcc | ||
|
|
ed40be066c | ||
|
|
fb092e7c25 | ||
|
|
55376b4752 | ||
|
|
113d663fce | ||
|
|
85aff0179a | ||
|
|
15d5a9197a | ||
|
|
e1bf06cb3d | ||
|
|
bbe7a4c217 | ||
|
|
ccedeba7ec | ||
|
|
fde3546574 | ||
|
|
7bcc64e0df | ||
|
|
b799485b74 | ||
|
|
db1de5eddd | ||
|
|
2992b7809d | ||
|
|
f7e0c9be2b | ||
|
|
5ef0ded757 | ||
|
|
4054449650 | ||
|
|
2737a4b1ef | ||
|
|
224fcc7803 | ||
|
|
5c5fe512c8 | ||
|
|
a7a79eb75b | ||
|
|
2a7b53a47b | ||
|
|
5b67143782 | ||
|
|
a17ecd02c1 | ||
|
|
6271d12426 | ||
|
|
45d23efcf1 | ||
|
|
4deda8542d | ||
|
|
4548beb723 | ||
|
|
4f62066c63 | ||
|
|
3e3a457fda | ||
|
|
091689f9f5 | ||
|
|
bd6c136608 | ||
|
|
af9ee4e671 | ||
|
|
6ef040b4b8 | ||
|
|
a5a1169829 | ||
|
|
9e337ed02a | ||
|
|
ab5e211019 | ||
|
|
760b2a0f48 | ||
|
|
778f4670be | ||
|
|
f485825bfb | ||
|
|
9e575ab48c | ||
|
|
2774ed921b | ||
|
|
4a2cb376da | ||
|
|
846b0b64b9 | ||
|
|
557d79c173 | ||
|
|
42aed8e9e3 | ||
|
|
9688e94b4c | ||
|
|
a7d842a51c | ||
|
|
4f273e9b21 | ||
|
|
8b8d1f0b5e | ||
|
|
0655af3713 | ||
|
|
f94eb86321 | ||
|
|
344640eb26 | ||
|
|
b137c998b3 | ||
|
|
f23fc5d966 | ||
|
|
927f1f1df7 | ||
|
|
a8c96a2c81 | ||
|
|
15e828b7a6 | ||
|
|
c3aa41d679 | ||
|
|
c50a7292c7 | ||
|
|
47d3f4d279 | ||
|
|
3fe2ee5075 | ||
|
|
4dbc553b04 | ||
|
|
8fa490a82a | ||
|
|
d4f13e40b6 | ||
|
|
86815913c1 | ||
|
|
85f93c437f | ||
|
|
940d2cc714 | ||
|
|
b8a5eaf9a6 | ||
|
|
4ba3fdbe35 | ||
|
|
33a66da2f1 | ||
|
|
eb75978da4 | ||
|
|
448227c218 | ||
|
|
1f6ddcf62a | ||
|
|
10bab420bb | ||
|
|
067320e9c5 | ||
|
|
3e925cc9ac | ||
|
|
ffc26709ad | ||
|
|
985bd15968 | ||
|
|
75fad23464 | ||
|
|
4f70b85ccc | ||
|
|
a4bb5dc2d7 | ||
|
|
8abeeaa953 | ||
|
|
da239bb7e1 | ||
|
|
2e3079c01b | ||
|
|
986663cd97 | ||
|
|
11d25afa56 | ||
|
|
3c8c0b9505 | ||
|
|
b8d6b9072d | ||
|
|
3033d088b4 | ||
|
|
c11d98fa17 | ||
|
|
ad5f03d475 | ||
|
|
ba36814894 | ||
|
|
56eed388a8 | ||
|
|
ffaf5fd3d8 | ||
|
|
2df2c63994 | ||
|
|
b5e6bec141 | ||
|
|
bd56e0420d | ||
|
|
d072245ab9 | ||
|
|
058aa98507 | ||
|
|
227bdf015f | ||
|
|
d074cd5968 | ||
|
|
a4373c0949 | ||
|
|
95661a5da1 | ||
|
|
e2a3df86e1 |
@@ -0,0 +1,66 @@
|
||||
# Coverage configuration for pytest-cov
|
||||
[run]
|
||||
source = .
|
||||
omit =
|
||||
*/tests/*
|
||||
*/venv/*
|
||||
*/env/*
|
||||
*/__pycache__/*
|
||||
*/migrations/*
|
||||
*/conftest.py
|
||||
*/setup.py
|
||||
*/manage.py
|
||||
*/.tox/*
|
||||
*/node_modules/*
|
||||
*/static/*
|
||||
*/media/*
|
||||
*/htmlcov/*
|
||||
*/.claude/*
|
||||
restart_gallery.py
|
||||
standalone_tagger.py
|
||||
|
||||
branch = true
|
||||
parallel = true
|
||||
|
||||
[report]
|
||||
# Regexes for lines to exclude from consideration
|
||||
exclude_lines =
|
||||
# Have to re-enable the standard pragma
|
||||
pragma: no cover
|
||||
|
||||
# Don't complain about missing debug-only code:
|
||||
def __repr__
|
||||
if self\.debug
|
||||
|
||||
# Don't complain if tests don't hit defensive assertion code:
|
||||
raise AssertionError
|
||||
raise NotImplementedError
|
||||
|
||||
# Don't complain if non-runnable code isn't run:
|
||||
if 0:
|
||||
if __name__ == .__main__.:
|
||||
|
||||
# Don't complain about abstract methods, they aren't run:
|
||||
@(abc\.)?abstractmethod
|
||||
|
||||
# Don't complain about type checking imports
|
||||
if TYPE_CHECKING:
|
||||
|
||||
# Don't complain about logger calls
|
||||
\.logger\.(debug|info|warning|error|critical)
|
||||
|
||||
ignore_errors = true
|
||||
show_missing = true
|
||||
skip_covered = false
|
||||
precision = 2
|
||||
|
||||
[html]
|
||||
directory = htmlcov
|
||||
title = ComfyUI PromptManager Test Coverage Report
|
||||
|
||||
[xml]
|
||||
output = coverage.xml
|
||||
|
||||
[json]
|
||||
output = coverage.json
|
||||
pretty_print = true
|
||||
@@ -0,0 +1,6 @@
|
||||
[flake8]
|
||||
max-line-length = 88
|
||||
extend-ignore = E203, W503
|
||||
exclude = reference/, .git, __pycache__, web/, venv/, htmlcov/, build/, dist/, .claude/, docs/, standalone_tagger.py
|
||||
per-file-ignores =
|
||||
tests/*:E402
|
||||
@@ -0,0 +1,29 @@
|
||||
{
|
||||
"required_status_checks": {
|
||||
"strict": true,
|
||||
"contexts": [
|
||||
"py3.10 on ubuntu-latest",
|
||||
"py3.11 on ubuntu-latest",
|
||||
"py3.12 on ubuntu-latest",
|
||||
"py3.13 on ubuntu-latest",
|
||||
"py3.10 on windows-latest",
|
||||
"py3.11 on windows-latest",
|
||||
"py3.12 on windows-latest",
|
||||
"py3.13 on windows-latest",
|
||||
"py3.10 on macos-latest",
|
||||
"py3.11 on macos-latest",
|
||||
"py3.12 on macos-latest",
|
||||
"py3.13 on macos-latest",
|
||||
"coverage gate (85%)",
|
||||
"black and flake8",
|
||||
"frontend modules",
|
||||
"viewerjs",
|
||||
"real ComfyUI compatibility"
|
||||
]
|
||||
},
|
||||
"enforce_admins": true,
|
||||
"required_pull_request_reviews": null,
|
||||
"restrictions": null,
|
||||
"allow_force_pushes": false,
|
||||
"allow_deletions": false
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
version: 2
|
||||
|
||||
updates:
|
||||
# Python dependencies (requirements.txt / setup.cfg / pyproject.toml)
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
day: "sunday"
|
||||
time: "06:00"
|
||||
timezone: "America/Los_Angeles"
|
||||
open-pull-requests-limit: 10
|
||||
labels:
|
||||
- "dependencies"
|
||||
- "security"
|
||||
commit-message:
|
||||
prefix: "deps(pip)"
|
||||
include: "scope"
|
||||
rebase-strategy: "auto"
|
||||
# Optional: group non-breaking updates to reduce PR noise
|
||||
groups:
|
||||
python-minor-and-patch:
|
||||
applies-to: "version-updates"
|
||||
update-types:
|
||||
- "minor"
|
||||
- "patch"
|
||||
|
||||
# GitHub Actions workflow dependencies
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
day: "sunday"
|
||||
time: "06:15"
|
||||
timezone: "America/Los_Angeles"
|
||||
open-pull-requests-limit: 10
|
||||
labels:
|
||||
- "dependencies"
|
||||
commit-message:
|
||||
prefix: "deps(actions)"
|
||||
include: "scope"
|
||||
rebase-strategy: "auto"
|
||||
groups:
|
||||
actions-minor-and-patch:
|
||||
applies-to: "version-updates"
|
||||
update-types:
|
||||
- "minor"
|
||||
- "patch"
|
||||
|
||||
# Maintainer browser test dependencies; Viewer.js itself has an asset-aware updater.
|
||||
- package-ecosystem: "npm"
|
||||
directory: "/tools/vendor"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
allow:
|
||||
- dependency-name: "@playwright/test"
|
||||
labels:
|
||||
- "dependencies"
|
||||
commit-message:
|
||||
prefix: "deps(vendor-tools)"
|
||||
@@ -0,0 +1,62 @@
|
||||
name: Code Quality
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
|
||||
jobs:
|
||||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
- name: Set up Python 3.10
|
||||
uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: "3.10"
|
||||
|
||||
- name: Cache pip dependencies
|
||||
uses: actions/cache@v6
|
||||
with:
|
||||
path: ~/.cache/pip
|
||||
key: ${{ runner.os }}-pip-quality-${{ hashFiles('requirements-dev.txt') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-pip-quality-
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
- name: Check formatting with Black
|
||||
run: black --check --diff .
|
||||
|
||||
- name: Lint with flake8
|
||||
run: |
|
||||
# Fail on syntax errors and undefined names
|
||||
flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
|
||||
# Warnings as non-blocking
|
||||
flake8 . --count --exit-zero --max-line-length=88 --statistics --exclude=reference/,.git,__pycache__,web/
|
||||
|
||||
security:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
- name: Set up Python 3.10
|
||||
uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: "3.10"
|
||||
|
||||
- name: Install bandit
|
||||
run: pip install bandit[toml]
|
||||
|
||||
- name: Security scan
|
||||
run: bandit -r . -ll --exclude=./tests,./reference,./venv -q || true
|
||||
@@ -18,7 +18,7 @@ jobs:
|
||||
if: ${{ github.repository_owner == 'ComfyAssets' }}
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v7
|
||||
with:
|
||||
submodules: true
|
||||
- name: Publish Custom Node
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
name: Tests
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches: [main]
|
||||
push:
|
||||
branches: [main]
|
||||
|
||||
jobs:
|
||||
test:
|
||||
name: py${{ matrix.python-version }} on ${{ matrix.os }}
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest, macos-latest]
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
cache: pip
|
||||
cache-dependency-path: requirements-dev.txt
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
- name: Run tests
|
||||
run: python -m unittest discover -s tests -v
|
||||
|
||||
comfyui:
|
||||
name: real ComfyUI compatibility
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 25
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
with:
|
||||
persist-credentials: false
|
||||
- uses: actions/setup-node@v7
|
||||
with:
|
||||
node-version: "24"
|
||||
cache: npm
|
||||
cache-dependency-path: tools/vendor/package-lock.json
|
||||
- name: Read the reviewed runtime pins
|
||||
id: runtime
|
||||
run: |
|
||||
node --input-type=module <<'JS' >> "$GITHUB_OUTPUT"
|
||||
import runtime from './tools/vendor/comfyui-runtime.json' with { type: 'json' };
|
||||
for (const [key, value] of Object.entries(runtime)) console.log(`${key}=${value}`);
|
||||
JS
|
||||
- name: Check out pinned ComfyUI
|
||||
uses: actions/checkout@v7
|
||||
with:
|
||||
repository: ${{ steps.runtime.outputs.repository }}
|
||||
ref: ${{ steps.runtime.outputs.revision }}
|
||||
path: .ci-comfyui
|
||||
persist-credentials: false
|
||||
- uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: ${{ steps.runtime.outputs.pythonVersion }}
|
||||
cache: pip
|
||||
cache-dependency-path: |
|
||||
requirements-dev.txt
|
||||
.ci-comfyui/requirements.txt
|
||||
tools/vendor/comfyui-runtime.json
|
||||
- name: Install CPU runtime and test dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install "torch==${{ steps.runtime.outputs.torchVersion }}" "torchvision==${{ steps.runtime.outputs.torchvisionVersion }}" --index-url https://download.pytorch.org/whl/cpu
|
||||
pip install -r .ci-comfyui/requirements.txt -r requirements-dev.txt "torch==${{ steps.runtime.outputs.torchVersion }}" "torchvision==${{ steps.runtime.outputs.torchvisionVersion }}"
|
||||
- run: npm ci --prefix tools/vendor --ignore-scripts
|
||||
- run: npx --prefix tools/vendor playwright install --with-deps chromium
|
||||
- name: Classic and Nodes 2.0 startup and compatibility
|
||||
env:
|
||||
PM_COMFY_ROOT: ${{ github.workspace }}/.ci-comfyui
|
||||
PM_TEST_PYTHON: python
|
||||
run: npm run --prefix tools/vendor test:comfyui
|
||||
- name: Upload real-host reports and startup logs
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: real-comfyui-reports
|
||||
path: tools/vendor/test-results/
|
||||
if-no-files-found: ignore
|
||||
retention-days: 14
|
||||
|
||||
coverage:
|
||||
name: coverage gate (85%)
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
- uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: "3.12"
|
||||
cache: pip
|
||||
cache-dependency-path: requirements-dev.txt
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
- name: Measure coverage
|
||||
run: |
|
||||
python -m coverage run -m unittest discover -s tests
|
||||
python -m coverage combine
|
||||
python -m coverage report --fail-under=85
|
||||
|
||||
lint:
|
||||
name: black and flake8
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
- uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: "3.12"
|
||||
- run: pip install black flake8
|
||||
- run: black --check .
|
||||
- run: flake8 .
|
||||
|
||||
js-test:
|
||||
name: frontend modules
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
- uses: actions/setup-node@v7
|
||||
with:
|
||||
node-version: "24"
|
||||
|
||||
- name: Check shipped extension and page modules parse
|
||||
run: |
|
||||
shopt -s globstar nullglob
|
||||
for f in web/comfy/**/*.js web/comfy/**/*.mjs web/js/**/*.js web/js/**/*.mjs; do
|
||||
node --input-type=module --check < "$f" || exit 1
|
||||
done
|
||||
|
||||
- name: Run JS tests
|
||||
run: node --test tests/js/*.test.js
|
||||
@@ -0,0 +1,47 @@
|
||||
name: Viewer.js vendor and compatibility
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches: [main]
|
||||
push:
|
||||
branches: [main]
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
viewerjs:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
- uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: '3.12'
|
||||
cache: pip
|
||||
cache-dependency-path: requirements-dev.txt
|
||||
- run: python -m pip install -r requirements-dev.txt
|
||||
- uses: actions/setup-node@v7
|
||||
with:
|
||||
node-version: '24'
|
||||
cache: npm
|
||||
cache-dependency-path: tools/vendor/package-lock.json
|
||||
- run: npm ci --prefix tools/vendor --ignore-scripts
|
||||
- run: npm run --prefix tools/vendor check
|
||||
- run: npm run --prefix tools/vendor test
|
||||
- run: npx --prefix tools/vendor playwright install --with-deps chromium firefox webkit
|
||||
- name: Frontend unit, viewer, workflow, and coverage checks
|
||||
run: npm run --prefix tools/vendor test:coverage
|
||||
- name: Firefox and WebKit page workflows
|
||||
run: npm run --prefix tools/vendor test:workflows -- --project firefox --project webkit
|
||||
- name: Upload browser traces and frontend coverage
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: frontend-test-reports
|
||||
path: |
|
||||
tools/vendor/test-results/
|
||||
coverage/js/
|
||||
!coverage/js/raw/**
|
||||
retention-days: 14
|
||||
@@ -0,0 +1,83 @@
|
||||
name: Update vendored Viewer.js
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: '30 7 * * *'
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
version:
|
||||
description: 'Stable version to vendor; leave blank for the latest npm release'
|
||||
required: false
|
||||
type: string
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: viewerjs-update
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
update:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
actions: write
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
with:
|
||||
ref: ${{ github.event.repository.default_branch }}
|
||||
- uses: actions/setup-node@v7
|
||||
with:
|
||||
node-version: '24'
|
||||
cache: npm
|
||||
cache-dependency-path: tools/vendor/package-lock.json
|
||||
- run: npm ci --prefix tools/vendor --ignore-scripts
|
||||
- name: Update the dependency and regenerate assets
|
||||
env:
|
||||
VIEWER_VERSION: ${{ inputs.version }}
|
||||
run: |
|
||||
if [ -n "$VIEWER_VERSION" ]; then
|
||||
npm run --prefix tools/vendor update -- --version "$VIEWER_VERSION"
|
||||
else
|
||||
npm run --prefix tools/vendor update
|
||||
fi
|
||||
- name: Detect changes
|
||||
id: changes
|
||||
run: |
|
||||
if git diff --quiet -- tools/vendor/package.json tools/vendor/package-lock.json web/lib/viewerjs; then
|
||||
echo 'changed=false' >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo 'changed=true' >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
- if: steps.changes.outputs.changed == 'true'
|
||||
run: |
|
||||
npm run --prefix tools/vendor check
|
||||
npm run --prefix tools/vendor test
|
||||
node --test tests/js/*.test.js
|
||||
npx --prefix tools/vendor playwright install --with-deps chromium
|
||||
npm run --prefix tools/vendor test:browser
|
||||
- name: Open the update PR
|
||||
if: steps.changes.outputs.changed == 'true'
|
||||
env:
|
||||
BASE_BRANCH: ${{ github.event.repository.default_branch }}
|
||||
run: |
|
||||
version=$(node -p "require('./web/lib/viewerjs/provenance.json').version")
|
||||
branch="deps/viewerjs-$version"
|
||||
existing=$(gh pr list --repo "$GITHUB_REPOSITORY" --state all --head "$branch" --json number --jq 'length')
|
||||
if [ "$existing" != '0' ]; then
|
||||
echo "A PR already exists for Viewer.js $version; leaving it for review."
|
||||
exit 0
|
||||
fi
|
||||
remote=$(git ls-remote origin "refs/heads/$branch")
|
||||
if [ -z "$remote" ]; then
|
||||
git switch -c "$branch"
|
||||
git add tools/vendor/package.json tools/vendor/package-lock.json web/lib/viewerjs
|
||||
git -c user.name='github-actions[bot]' -c user.email='41898282+github-actions[bot]@users.noreply.github.com' commit -m "deps: update vendored Viewer.js to $version"
|
||||
git push origin "HEAD:refs/heads/$branch"
|
||||
fi
|
||||
gh pr create --repo "$GITHUB_REPOSITORY" --base "$BASE_BRANCH" --head "$branch" --title "deps: update vendored Viewer.js to $version" --body-file tools/vendor/update-pr.md
|
||||
gh workflow run viewerjs-check.yml --repo "$GITHUB_REPOSITORY" --ref "$branch"
|
||||
@@ -1,3 +1,29 @@
|
||||
# Tailwind standalone CLI (downloaded by `make css-setup`)
|
||||
tailwindcss
|
||||
|
||||
# AI tooling / local config
|
||||
.claude/
|
||||
.playwright-mcp/
|
||||
.serena/
|
||||
AGENTS.md
|
||||
CLAUDE.md
|
||||
|
||||
# Local files
|
||||
# Local PromptManager state (config.json holds API keys and machine paths)
|
||||
config.json
|
||||
resume.sh
|
||||
*.backup_*
|
||||
promptmanager_logs/
|
||||
custom_data/
|
||||
data/
|
||||
docs/
|
||||
logs/
|
||||
standalone_tagger.py
|
||||
prompts.db
|
||||
prompts.db-shm
|
||||
prompts.db-wal
|
||||
prompts.db-journal
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
@@ -16,6 +42,8 @@ eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
!web/lib/
|
||||
!web/lib/**
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
@@ -172,3 +200,14 @@ cython_debug/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
# Maintainer-only JavaScript tooling and browser test artifacts
|
||||
node_modules/
|
||||
test-results/
|
||||
playwright-report/
|
||||
|
||||
# Combined frontend coverage reports
|
||||
coverage/
|
||||
|
||||
# Disposable upstream checkout for real-host CI
|
||||
.ci-comfyui/
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
repos:
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: black
|
||||
name: black
|
||||
entry: black --line-length=88 --target-version=py310
|
||||
language: system
|
||||
types: [python]
|
||||
|
||||
- id: flake8
|
||||
name: flake8
|
||||
entry: flake8 --show-source
|
||||
language: system
|
||||
types: [python]
|
||||
|
||||
- id: bandit
|
||||
name: bandit
|
||||
entry: bash -c 'bandit "$@" -ll -q 2>&1; true'
|
||||
language: system
|
||||
types: [python]
|
||||
exclude: ^(tests/|reference/)
|
||||
|
||||
- id: tests
|
||||
name: tests
|
||||
entry: python -m unittest discover tests/ -v
|
||||
language: system
|
||||
pass_filenames: false
|
||||
stages: [pre-push]
|
||||
@@ -0,0 +1,43 @@
|
||||
TAILWIND_VERSION := 4.1.18
|
||||
TAILWIND_CLI := ./tailwindcss
|
||||
UNAME_S := $(shell uname -s)
|
||||
UNAME_M := $(shell uname -m)
|
||||
|
||||
ifeq ($(UNAME_S),Darwin)
|
||||
ifeq ($(UNAME_M),arm64)
|
||||
TAILWIND_PLATFORM := tailwindcss-macos-arm64
|
||||
else
|
||||
TAILWIND_PLATFORM := tailwindcss-macos-x64
|
||||
endif
|
||||
else
|
||||
ifeq ($(UNAME_M),aarch64)
|
||||
TAILWIND_PLATFORM := tailwindcss-linux-arm64
|
||||
else
|
||||
TAILWIND_PLATFORM := tailwindcss-linux-x64
|
||||
endif
|
||||
endif
|
||||
|
||||
TAILWIND_URL := https://github.com/tailwindlabs/tailwindcss/releases/download/v$(TAILWIND_VERSION)/$(TAILWIND_PLATFORM)
|
||||
|
||||
.PHONY: css css-watch css-setup clean-css
|
||||
|
||||
## Download Tailwind standalone CLI (dev only, not committed)
|
||||
css-setup:
|
||||
@if [ ! -f $(TAILWIND_CLI) ]; then \
|
||||
echo "Downloading Tailwind CSS v$(TAILWIND_VERSION) CLI..."; \
|
||||
curl -sLo $(TAILWIND_CLI) $(TAILWIND_URL) && chmod +x $(TAILWIND_CLI); \
|
||||
else \
|
||||
echo "Tailwind CLI already present"; \
|
||||
fi
|
||||
|
||||
## Rebuild compiled CSS (run after changing Tailwind classes in HTML/JS)
|
||||
css: css-setup
|
||||
$(TAILWIND_CLI) -i web/lib/tailwind/input.css -o web/lib/tailwind/styles.css --minify
|
||||
|
||||
## Watch mode for development
|
||||
css-watch: css-setup
|
||||
$(TAILWIND_CLI) -i web/lib/tailwind/input.css -o web/lib/tailwind/styles.css --watch
|
||||
|
||||
## Remove downloaded CLI binary
|
||||
clean-css:
|
||||
rm -f $(TAILWIND_CLI)
|
||||
@@ -1,48 +1,12 @@
|
||||
# ComfyUI Prompt Manager
|
||||
|
||||
A comprehensive ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, automatic image gallery system, and powerful ComfyUI workflow metadata analysis using SQLite.
|
||||

|
||||

|
||||

|
||||

|
||||

|
||||
|
||||
# 🚀 Prompt Manager v2 Rewrite
|
||||
|
||||
We’ve been hard at work on a **complete rewrite** of Prompt Manager.
|
||||
The new version brings over all features from v1, introduces a dark theme (v2 look), and adds **negative prompt support**.
|
||||
|
||||
## ✨ What’s New
|
||||
- Unified dark theme UI (v2)
|
||||
- Table view (for power users) and Card view (for quick browsing)
|
||||
- Negative prompt field with clear highlighting
|
||||
- Streamlined toolbar & filters
|
||||
|
||||
---
|
||||
|
||||
## 📊 Progress
|
||||
|
||||
**Rewrite completion status:**
|
||||
|
||||
- [█████████░░] 85% done
|
||||
|
||||
✅ Core features migrated
|
||||
✅ Dark UI overhaul
|
||||
✅ Table + Card views
|
||||
✅ Negative prompt handling
|
||||
⬜ API integration & real data wiring
|
||||
⬜ Community feedback + polish pass
|
||||
|
||||
---
|
||||
|
||||
## 🧪 Call for Beta Testers
|
||||
|
||||
We’d love your help testing Prompt Manager v2 before release!
|
||||
|
||||
- Try the new UI locally
|
||||
- Report bugs, UI quirks, or feature requests
|
||||
- Share feedback on usability & performance
|
||||
|
||||
👉 Open an [Issue](../../issues) with feedback or join the discussion in [Discussions](../../discussions).
|
||||
|
||||
---
|
||||
|
||||
💡 Thanks to everyone who has supported this project—your input drives its evolution!
|
||||
A comprehensive ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, automatic image gallery system, folder-based organization, LoRA Manager integration, and powerful ComfyUI workflow metadata analysis using SQLite.
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -67,12 +31,24 @@ A text-focused variant that outputs raw STRING for maximum flexibility:
|
||||
- **🔗 Text Concatenation**: Built-in prepend/append functionality for dynamic text assembly
|
||||
- **⚡ Lightweight Processing**: No CLIP encoding overhead for text-only workflows
|
||||
|
||||
### **PromptSearchList** (Batch Processing Node)
|
||||
|
||||
A search node that outputs prompts as a list for batch processing workflows:
|
||||
|
||||
- **🔍 Database Search**: Search prompts by text, category, tags, and rating
|
||||
- **📋 List Output**: Uses `OUTPUT_IS_LIST` for native batch processing support
|
||||
- **🔗 Batch Compatible**: Connect directly to batch processing nodes like String OutputList
|
||||
- **⚡ Read-Only**: Lightweight search operation without database writes
|
||||
|
||||
This node enables powerful batch workflows by allowing you to search your prompt database and process multiple prompts in a single execution.
|
||||
|
||||
### **Shared Features**
|
||||
|
||||
Both nodes include the complete PromptManager feature set:
|
||||
|
||||
- **💾 Persistent Storage**: Automatically saves all prompts to a local SQLite database
|
||||
- **🔍 Advanced Search**: Query past prompts with text search, category filtering, and metadata
|
||||
- **📁 Folder Filter**: Browse and filter prompts by output subdirectory
|
||||
- **🖼️ Automatic Image Gallery**: Automatically links generated images to their prompts
|
||||
- **🏷️ Rich Metadata**: Add categories, tags, ratings, notes, and workflow names to prompts
|
||||
- **🚫 Duplicate Prevention**: Uses SHA256 hashing to detect and prevent duplicate storage
|
||||
@@ -81,6 +57,8 @@ Both nodes include the complete PromptManager feature set:
|
||||
- **🔬 Workflow Analysis**: Extract and analyze ComfyUI workflow data from PNG images
|
||||
- **📋 Metadata Viewer**: Standalone tool for analyzing ComfyUI-generated images
|
||||
- **🛠️ System Management**: Built-in diagnostics, backup/restore, and maintenance tools
|
||||
- **🏷️ AI AutoTag**: Automatically tag your image collection using WD14 or JoyCaption vision models
|
||||
- **🔗 LoRA Manager Integration**: Import LoRA metadata and preview images from [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manager)
|
||||
|
||||

|
||||
|
||||
@@ -124,6 +102,7 @@ The comprehensive gallery system provides complete management of all ComfyUI out
|
||||
- **Performance Optimized**: Efficient scanning with thumbnail generation and lazy loading
|
||||
|
||||
#### **🎬 Full Video Support**
|
||||
|
||||
- **Native Video Playback**: HTML5 video player with custom controls for all supported formats
|
||||
- **Video Thumbnails**: Automatic thumbnail generation with OpenCV, ffmpeg, or placeholder fallback
|
||||
- **Playback Controls**: Autoplay, mute, loop settings with persistent preferences
|
||||
@@ -132,6 +111,7 @@ The comprehensive gallery system provides complete management of all ComfyUI out
|
||||
- **Mixed Media Gallery**: Videos and images integrated seamlessly in the same interface
|
||||
|
||||
#### **🔍 Professional Image Viewer (ViewerJS Integration)**
|
||||
|
||||
- **Full-Featured Viewing**: Zoom, rotation, slideshow modes with professional controls
|
||||
- **Advanced Navigation**: Smooth transitions, thumbnail strip, and keyboard shortcuts
|
||||
- **Zoom Controls**: Mouse wheel zoom, drag to pan, and touch gesture support
|
||||
@@ -150,6 +130,7 @@ The comprehensive gallery system provides complete management of all ComfyUI out
|
||||
- **Cross-Platform Support**: Works on desktop, tablet, and mobile devices
|
||||
|
||||
#### **⚡ Performance Features**
|
||||
|
||||
- **Smart Thumbnail System**: Generates optimized thumbnails for faster loading
|
||||
- **Lazy Loading**: Load content as needed for smooth browsing of large galleries
|
||||
- **Grid & List Views**: Choose between compact grid or detailed list display
|
||||
@@ -158,6 +139,29 @@ The comprehensive gallery system provides complete management of all ComfyUI out
|
||||
|
||||
_Powered by [ViewerJS](https://github.com/fengyuanchen/viewerjs) for professional image viewing capabilities._
|
||||
|
||||
Viewer.js is vendored locally. See [Viewer.js maintenance](tools/vendor/README.md)
|
||||
for reproducible updates, verification, and automated update PRs.
|
||||
|
||||
Prompt Manager appears as **Prompt Manager (Classic)** in the Add Node menu.
|
||||
It keeps the existing layout and behavior while using shared
|
||||
[ComfyAssets Toolkit](https://github.com/ComfyAssets/ComfyAssets-Toolkit) controls
|
||||
that follow ComfyUI's light/dark theme. Existing workflows keep the same internal
|
||||
node identifier and sockets. The reusable architecture is inspired by
|
||||
[Pixaroma](https://github.com/pixaroma/ComfyUI-Pixaroma); no implementation or assets
|
||||
were copied. See the [toolkit maintenance notes](tools/vendor/README.md#shared-comfyassets-browser-controls).
|
||||
|
||||
**Prompt Manager (Focused)** adds an editor-first layout with collapsible
|
||||
composition and metadata fields. Open **Browse library** to search saved prompts,
|
||||
filter by saved category and comma-separated tags, preview a selection, then choose
|
||||
**Use prompt**; **Load metadata** controls whether
|
||||
category and tags come with it. Metadata fields suggest saved categories and tags;
|
||||
choose with the mouse or arrow keys and Enter. Tag selections append `, ` for the
|
||||
next tag; category remains a single value. New values can still be typed freely.
|
||||
Prepend and append STRING sockets remain visible
|
||||
when the fields are collapsed, so LoRA trigger outputs can connect directly.
|
||||
Encoding, storage, and CONDITIONING/text outputs match Classic. Both layouts
|
||||
support ComfyUI Classic and Nodes 2.0.
|
||||
|
||||
### 🔬 ComfyUI Workflow Analysis
|
||||
|
||||
- **PNG Metadata Parsing**: Advanced PNG chunk analysis to extract embedded ComfyUI data
|
||||
@@ -183,7 +187,7 @@ _Powered by [ViewerJS](https://github.com/fengyuanchen/viewerjs) for professiona
|
||||
|
||||
- **Full-text search** across all stored prompts
|
||||
- **Category filtering** for organized prompt collections
|
||||
- **Tag-based search** with support for multiple tags
|
||||
- **Tag-based search** with support for multiple tags, with autocomplete that narrows to tags seen together with the ones already typed
|
||||
- **Rating filters** to find your best prompts
|
||||
- **Date range filtering** for temporal searches
|
||||
- **Recent prompts** quick access
|
||||
@@ -279,6 +283,20 @@ For text-only workflows or when you need STRING output:
|
||||
- **Append Text**: Text added after your main prompt
|
||||
4. **Connect the STRING output** to any node that accepts text input
|
||||
|
||||
#### Using PromptSearchList (Batch Processing)
|
||||
|
||||
For batch processing workflows that need multiple prompts:
|
||||
|
||||
1. **Add the PromptSearchList node** to your workflow
|
||||
2. **Configure search filters**:
|
||||
- **text**: Search for prompts containing this text
|
||||
- **category**: Filter by specific category
|
||||
- **tags**: Comma-separated tags to filter by (e.g., "anime, detailed")
|
||||
- **min_rating**: Minimum rating (0-5) to include
|
||||
- **limit**: Maximum number of results (1-1000, default 50)
|
||||
3. **Connect the prompts output** to batch-compatible nodes like String OutputList
|
||||
4. **Each matching prompt** will be processed as a separate batch item
|
||||
|
||||
#### Shared Metadata Options
|
||||
|
||||
Both nodes support the same metadata fields:
|
||||
@@ -296,22 +314,30 @@ Both nodes will automatically save prompts to the database and link any generate
|
||||
The comprehensive gallery system provides complete access to all your ComfyUI output:
|
||||
|
||||
#### **📁 Accessing the Gallery**
|
||||
|
||||
1. **Open the admin interface** at `http://localhost:8188/prompt_manager/admin`
|
||||
2. **Click the "🖼️ Gallery" button** to open the output-wide gallery
|
||||
3. **Browse all content** from your ComfyUI output directory
|
||||
|
||||
#### **🎬 Video & Image Support**
|
||||
|
||||
- **Mixed Media Viewing**: Images and videos appear together in the same gallery
|
||||
- **Video Thumbnails**: Videos show with play button overlays and 🎬 icons
|
||||
- **Smart Filtering**: Filter by media type (All, Images, Videos) using the dropdown
|
||||
|
||||
#### **⚙️ Gallery Settings**
|
||||
|
||||
Configure your viewing experience in the settings panel:
|
||||
- **Performance Settings**: Grid columns, image quality, lazy loading
|
||||
|
||||
- **Display Settings**: Default view, items per page, grid columns, media info, infinite scrolling
|
||||
- **Performance Settings**: Image quality, lazy loading
|
||||
- **Video Settings**: Autoplay, mute by default, loop videos
|
||||
- **Thumbnail Management**: Generate optimized thumbnails for faster loading
|
||||
- **Thumbnail Management**: Generate optimized thumbnails for faster loading (runs on the configured worker threads; a run that stops at the per-run cap tells you how many remain)
|
||||
- **Find Duplicate Images**: Hashes every file on the worker threads with a progress bar; close the window and it keeps running, results wait in a pill
|
||||
- **Rescan Output Folder**: Drops the server-side listing cache and rewalks the folders
|
||||
|
||||
#### **🖼️ Image Viewing Experience**
|
||||
|
||||
1. **Click any image** to open the professional viewer with:
|
||||
- **Full-screen viewing** with zoom, rotation, and slideshow controls
|
||||
- **Integrated metadata panel** showing workflow data and generation parameters
|
||||
@@ -323,6 +349,7 @@ Configure your viewing experience in the settings panel:
|
||||
- **Space** for slideshow mode
|
||||
|
||||
#### **🎬 Video Viewing Experience**
|
||||
|
||||
1. **Click any video** to open the video player with:
|
||||
- **HTML5 video controls** with play, pause, seek, and fullscreen
|
||||
- **Custom control buttons** for mute/unmute and loop toggle
|
||||
@@ -331,6 +358,7 @@ Configure your viewing experience in the settings panel:
|
||||
2. **Playback settings** are remembered across sessions
|
||||
|
||||
#### **🔍 Advanced Features**
|
||||
|
||||
- **Real-time Search**: Search across filenames and metadata
|
||||
- **View Modes**: Switch between grid and list views
|
||||
- **Sorting Options**: Sort by date, size, or filename
|
||||
@@ -356,14 +384,119 @@ Use the standalone metadata viewer to analyze any ComfyUI-generated PNG:
|
||||
Import existing ComfyUI images into your database:
|
||||
|
||||
1. **Open admin dashboard** and click **"📸 Scan Images"**
|
||||
2. **Start scan** to analyze your entire output directory
|
||||
3. **Monitor progress** with real-time status updates
|
||||
2. **Start scan** to analyze every configured output directory
|
||||
3. **Monitor progress** with real-time status updates. The scan runs on the server: closing the window only minimizes it to a progress pill, and reloading the page re-attaches to a running scan
|
||||
4. **Review results** showing:
|
||||
- Images found and processed
|
||||
- Prompts extracted and linked
|
||||
- Any errors or issues encountered
|
||||
5. **Access imported data** through the normal gallery interface
|
||||
|
||||
The scan reads PNG text chunks and A1111-style parameters, and, when `ffprobe` is installed, the prompt ComfyUI's video savers embed in MP4/MOV/WebM files. Files are read on the number of **worker threads** set in Settings → Performance (default: half your cores), which makes a large library several times faster.
|
||||
|
||||
### 🏷️ AI AutoTag
|
||||
|
||||
Automatically tag your entire image collection using AI vision models:
|
||||
|
||||
#### **Model Options**
|
||||
|
||||
- **WD14 SwinV2 / WD14 ViT**: Fast ONNX-based classifiers (~400MB) that output Danbooru tags with confidence scores — no prompt needed, adjustable confidence thresholds
|
||||
- **JoyCaption Beta One FP16**: Full precision model for highest quality tagging (requires more VRAM)
|
||||
- **JoyCaption Beta One GGUF (FP8)**: Quantized model for lower VRAM usage with minimal quality loss
|
||||
|
||||
#### **Two Tagging Modes**
|
||||
|
||||
1. **AutoTag (Batch Mode)**: Tag your entire collection automatically
|
||||
- Click **"🏷️ AutoTag"** in the admin dashboard
|
||||
- Choose your model type (FP16 or GGUF)
|
||||
- Select how to handle already-tagged images:
|
||||
- **Skip images with existing tags**: Preserve your meticulous manual tagging work
|
||||
- **Re-tag all images**: Overwrite existing tags with fresh AI analysis
|
||||
- Customize the system prompt to match your collection style
|
||||
- Monitor real-time progress as images are processed
|
||||
|
||||
2. **Review Mode**: Tag images one-by-one with approval
|
||||
- Review each image and its AI-generated tags before applying
|
||||
- Edit, add, or remove tags before saving
|
||||
- Skip images you don't want to tag
|
||||
- Perfect for curating high-quality tag assignments
|
||||
|
||||
#### **Handling Existing Tags**
|
||||
|
||||
For users who have meticulously tagged their collections from the start:
|
||||
- AutoTag respects your existing work with the **"Skip tagged"** option
|
||||
- **"auto-scanned"** placeholder tags don't count as real tags
|
||||
- Choose to re-tag specific images while preserving others
|
||||
- Review mode asks before overwriting on each image
|
||||
|
||||
#### **Customization**
|
||||
|
||||
Adjust the system prompt to match your tagging style:
|
||||
- Focus on specific attributes (style, mood, composition, subjects)
|
||||
- Match your existing tag vocabulary
|
||||
- Optimize for your collection's theme (anime, photography, landscapes, etc.)
|
||||
|
||||
#### **Requirements**
|
||||
|
||||
- Sufficient VRAM for the chosen model (FP16 requires more, GGUF is lighter)
|
||||
- Models are downloaded automatically on first use
|
||||
- Feature requests and improvements are welcome!
|
||||
|
||||
### 📁 Folder Filter
|
||||
|
||||

|
||||
|
||||
Organize and filter your prompt library by output subdirectory. The dashboard search panel includes a **Folder** dropdown that lists all subdirectories where your images are stored.
|
||||
|
||||
- **Automatic Detection**: Folders are extracted from your image paths — no configuration needed
|
||||
- **Multi-Directory Support**: Works across all configured gallery scan directories
|
||||
- **Quick Filtering**: Select a folder from the dropdown to instantly filter prompts to that subdirectory
|
||||
|
||||
> **Note:** If you have an existing prompt library from before the folder feature was added, you will need to **rescan your image library** (click **Scan Images** in the admin dashboard) to populate folder data for your existing prompts.
|
||||
|
||||
### 🔗 LoRA Manager Integration
|
||||
|
||||

|
||||
|
||||
If you use [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manager), PromptManager can import your LoRA metadata, trigger words, and example images directly into your prompt database. This lets you search, tag, and browse your LoRA collection alongside your regular prompts.
|
||||
|
||||
> **WIP:** LoRA Manager support is a work in progress — this was a highly requested feature. Please [open issues](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues) for any bugs or feature requests.
|
||||
|
||||
#### **Enabling the Integration**
|
||||
|
||||
1. Open the admin dashboard and click **Settings**
|
||||
2. Scroll to the **Integrations** section — PromptManager will auto-detect if LoRA Manager is installed
|
||||
|
||||

|
||||
|
||||
3. Toggle the **LoraManager** switch to enable
|
||||
4. Optionally enable **Auto-inject trigger words** to automatically append LoRA trigger words when `<lora:name:weight>` is detected in your prompts
|
||||
5. Click **Save Settings** (requires a ComfyUI restart to take effect)
|
||||
|
||||
#### **Importing LoRA Data**
|
||||
|
||||
Click **Import LoRA Data** in the settings panel to start the import. A progress popup shows real-time status as each LoRA is processed:
|
||||
|
||||
- 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
|
||||
|
||||

|
||||
|
||||
After import, filter by the `lora-manager` category to browse your LoRA collection.
|
||||
|
||||
#### **CivitAI API Key (Optional)**
|
||||
|
||||
A CivitAI API key is **required if your library contains NSFW LoRAs** — CivitAI blocks unauthenticated access to NSFW preview images. Without a key, only SFW preview images are downloaded.
|
||||
|
||||
To add your key:
|
||||
1. Get your API key from [civitai.com/user/account](https://civitai.com/user/account)
|
||||
2. Paste it into the **CivitAI API Key** field in the Integrations settings
|
||||
3. Save and re-import to download any previously skipped images
|
||||
|
||||
> **Note:** Re-importing is safe — PromptManager skips LoRAs that were already imported. To do a fresh import, the previous `lora-manager` data is cleared automatically before re-scanning.
|
||||
|
||||
### 🌐 Web Interface Features
|
||||
|
||||
The comprehensive web interface provides:
|
||||
@@ -471,9 +604,29 @@ PromptManagerText → Text Processor → Style Applicator → Final Text Node
|
||||
"cyberpunk city" → style processing → "neon-lit cyberpunk metropolis at night"
|
||||
```
|
||||
|
||||
### PromptSearchList (Batch Processing) Examples
|
||||
|
||||
#### Example 6: Batch Processing with Tags
|
||||
|
||||
```
|
||||
Tags: "anime, detailed"
|
||||
Limit: 10
|
||||
→ Outputs: List of up to 10 prompts with both "anime" and "detailed" tags
|
||||
→ Each prompt is processed as a separate batch item
|
||||
```
|
||||
|
||||
#### Example 7: Category-Based Batch
|
||||
|
||||
```
|
||||
Category: "portraits"
|
||||
Min Rating: 4
|
||||
Limit: 20
|
||||
→ Outputs: List of up to 20 highly-rated portrait prompts for batch generation
|
||||
```
|
||||
|
||||
### Traditional Examples
|
||||
|
||||
#### Example 6: Abstract Art
|
||||
#### Example 8: Abstract Art
|
||||
|
||||
```
|
||||
Input: "Swirling colors in an abstract geometric pattern"
|
||||
@@ -524,10 +677,12 @@ CREATE TABLE generated_images (
|
||||
|
||||
- **`prompt_manager.py`** - Main ComfyUI node implementation (CLIP encoding)
|
||||
- **`prompt_manager_text.py`** - Text-only node implementation (STRING output)
|
||||
- **`prompt_search_list.py`** - Batch search node implementation (LIST output)
|
||||
- **`database/models.py`** - Database schema and connection management
|
||||
- **`database/operations.py`** - CRUD operations and search functionality
|
||||
- **`py/api.py`** - Web API endpoints for the interface
|
||||
- **`py/api/`** - Web API route modules (prompts, images, tags, lora integration, etc.)
|
||||
- **`py/config.py`** - Configuration management
|
||||
- **`py/lora_utils.py`** - LoRA Manager integration utilities
|
||||
- **`utils/hashing.py`** - SHA256 hashing for deduplication
|
||||
- **`utils/validators.py`** - Input validation and sanitization
|
||||
- **`utils/image_monitor.py`** - Automatic image detection system
|
||||
@@ -537,7 +692,9 @@ CREATE TABLE generated_images (
|
||||
- **`utils/diagnostics.py`** - System diagnostics and health checks
|
||||
- **`web/admin.html`** - Advanced admin dashboard with metadata panel
|
||||
- **`web/index.html`** - Simple web interface
|
||||
- **`web/prompt_manager.js`** - JavaScript functionality
|
||||
- **`web/comfy/prompt_manager.js`** - ComfyUI canvas extension (the only file ComfyUI loads)
|
||||
- **`web/js/admin.js`**, **`web/js/gallery.js`** - Dashboard and gallery JavaScript
|
||||
- **`web/js/tags-page.js`** - Tag management JavaScript
|
||||
- **`web/metadata.html`** - Standalone PNG metadata viewer
|
||||
|
||||
### File Structure
|
||||
@@ -547,14 +704,16 @@ ComfyUI_PromptManager/
|
||||
├── __init__.py # Node registration
|
||||
├── prompt_manager.py # Main node implementation (CLIP encoding)
|
||||
├── prompt_manager_text.py # Text-only node implementation (STRING output)
|
||||
├── prompt_search_list.py # Batch search node implementation (LIST output)
|
||||
├── database/
|
||||
│ ├── __init__.py
|
||||
│ ├── models.py # Database schema
|
||||
│ └── operations.py # Database operations
|
||||
├── py/
|
||||
│ ├── __init__.py
|
||||
│ ├── api.py # Web API endpoints
|
||||
│ └── config.py # Configuration
|
||||
│ ├── api/ # Web API route modules
|
||||
│ ├── config.py # Configuration
|
||||
│ └── lora_utils.py # LoRA Manager integration
|
||||
├── utils/
|
||||
│ ├── __init__.py
|
||||
│ ├── hashing.py # Hashing utilities
|
||||
@@ -566,9 +725,15 @@ ComfyUI_PromptManager/
|
||||
│ └── diagnostics.py # System diagnostics
|
||||
├── web/
|
||||
│ ├── admin.html # Advanced admin dashboard
|
||||
│ ├── gallery.html # Image gallery
|
||||
│ ├── index.html # Simple web interface
|
||||
│ ├── metadata.html # Standalone metadata viewer
|
||||
│ └── prompt_manager.js # JavaScript functionality
|
||||
│ ├── comfy/
|
||||
│ │ └── prompt_manager.js # ComfyUI canvas extension (WEB_DIRECTORY)
|
||||
│ └── js/
|
||||
│ ├── admin.js # Dashboard JavaScript
|
||||
│ ├── gallery.js # Gallery JavaScript
|
||||
│ └── tags-page.js # Tag management JavaScript
|
||||
├── tests/
|
||||
│ ├── __init__.py
|
||||
│ └── test_basic.py # Test suite
|
||||
@@ -582,11 +747,24 @@ ComfyUI_PromptManager/
|
||||
|
||||
### Database Settings
|
||||
|
||||
You can customize the database path by modifying the configuration:
|
||||
You can customize the database path by creating a `config.json` file in the extension root:
|
||||
|
||||
```python
|
||||
# In py/config.py
|
||||
DATABASE_PATH = "custom_path/prompts.db"
|
||||
```json
|
||||
{
|
||||
"database": {
|
||||
"default_path": "/path/to/your/prompts.db"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Relative paths are resolved against the extension directory. Absolute paths are used as-is, which is recommended for Docker deployments where the working directory may not persist:
|
||||
|
||||
```json
|
||||
{
|
||||
"database": {
|
||||
"default_path": "/data/custom_nodes/ComfyUI_PromptManager/prompts.db"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Gallery & Monitoring Settings
|
||||
@@ -617,11 +795,15 @@ Configure the web interface behavior:
|
||||
# Web UI settings (PromptManagerConfig class)
|
||||
RESULT_TIMEOUT = 5 # Seconds to auto-hide results in ComfyUI node
|
||||
WEBUI_DISPLAY_MODE = 'newtab' # 'popup' or 'newtab' for Web UI button
|
||||
INFINITE_SCROLL = False # Dashboard list loads the next page on scroll by default
|
||||
SHOW_TEST_BUTTON = False # Show API test button in node UI
|
||||
```
|
||||
|
||||
All of these are editable in the dashboard Settings modal and saved to `config.json` (which is private to the installation and holds the CivitAI key, so it is never committed).
|
||||
|
||||
### Performance Tuning
|
||||
|
||||
- **Worker threads** (Settings → Performance): threads used by the image rescan, thumbnail generation and duplicate detection. Default is half the detected cores, maximum is the core count. Pillow releases the GIL while decoding, so the scan and thumbnails scale almost linearly; keep a few cores free while ComfyUI is generating.
|
||||
- The database automatically creates indexes for optimal search performance
|
||||
- Regular `VACUUM` operations keep the database optimized
|
||||
- Consider backing up the database periodically
|
||||
@@ -664,10 +846,15 @@ db.model.backup_database("backup_prompts.db")
|
||||
### Running Tests
|
||||
|
||||
```bash
|
||||
cd KikoTextEncode
|
||||
python -m pytest tests/ -v
|
||||
cd ComfyUI_PromptManager
|
||||
pip install -r requirements-dev.txt
|
||||
python -m unittest discover -s tests # Python (what CI runs)
|
||||
node --test tests/js/*.test.js # frontend pure modules
|
||||
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.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
|
||||
|
||||
The project follows PEP 8 guidelines with:
|
||||
@@ -709,7 +896,7 @@ The project follows PEP 8 guidelines with:
|
||||
For debugging, you can enable verbose logging in the node:
|
||||
|
||||
```python
|
||||
# Add to kiko_text_encode.py
|
||||
# Add to prompt_manager.py
|
||||
import logging
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
```
|
||||
@@ -726,15 +913,19 @@ MIT License - see LICENSE file for details.
|
||||
|
||||
## Roadmap
|
||||
|
||||
### Completed in v3.0.0
|
||||
### Recently Completed
|
||||
|
||||
- **✅ PNG Metadata Analysis**: Complete ComfyUI workflow extraction from images
|
||||
- **✅ Standalone Metadata Viewer**: Dedicated tool for analyzing any ComfyUI image
|
||||
- **✅ Advanced Admin Dashboard**: Comprehensive management interface with modern UI
|
||||
- **✅ Integrated Metadata Panel**: Real-time workflow analysis in image viewer
|
||||
- **✅ Bulk Image Scanning**: Mass import of existing ComfyUI images
|
||||
- **✅ System Management Tools**: Backup, restore, diagnostics, and maintenance
|
||||
- **✅ Enhanced Error Handling**: Robust PNG parsing with NaN value cleaning
|
||||
- **✅ Security & Reliability Hardening (3.3.0)**: validated gallery roots, contained file serving, safe backups and restores, background scans, 85%+ test coverage on all three operating systems
|
||||
- **✅ Worker Threads**: scan, thumbnails and duplicate detection run on a configurable thread pool
|
||||
- **✅ Tag Autocomplete & Infinite Scrolling**: faster browsing in the dashboard and gallery
|
||||
- **✅ Queue-Time Run Counting (3.2.4)**: run counts and image linking follow the ComfyUI queue
|
||||
- **✅ LoRA Manager Integration**: Import LoRA metadata, trigger words, and preview images
|
||||
- **✅ Folder Filter**: Browse and filter prompts by output subdirectory
|
||||
- **✅ Multi-Directory Gallery**: Scan multiple output directories simultaneously
|
||||
- **✅ WD14 Tagger**: Fast ONNX-based auto-tagging with Danbooru tags
|
||||
- **✅ Tailwind v4 Migration**: ComfyUI theme token system for consistent styling
|
||||
- **✅ Tag Management Page**: Dedicated page for tag search, rename, merge, and delete
|
||||
- **✅ Junction Table Tags**: Normalized tag storage with proper foreign keys
|
||||
|
||||
### Planned Features
|
||||
|
||||
@@ -742,20 +933,108 @@ MIT License - see LICENSE file for details.
|
||||
- **🤝 Collaboration**: Share prompt collections with other users
|
||||
- **🧠 AI Suggestions**: Recommend similar prompts based on metadata analysis
|
||||
- **📈 Advanced Analytics**: Detailed usage statistics and trends with workflow insights
|
||||
- **🔌 Plugin System**: Support for third-party extensions and custom analyzers
|
||||
- **🎨 Enhanced Batch Processing**: Advanced bulk operations with metadata editing
|
||||
- **🔄 Workflow Templates**: Save and reuse common workflow patterns
|
||||
- **📊 Visual Analytics**: Charts and graphs for prompt effectiveness analysis
|
||||
|
||||
### Integration Ideas
|
||||
|
||||
- **Auto-tagging**: Use AI to automatically categorize prompts
|
||||
- **Workflow linking**: Connect prompts to specific workflow templates
|
||||
- **Image analysis**: Analyze generated images to improve suggestions
|
||||
- **Version control**: Track prompt iterations and effectiveness
|
||||
|
||||
## Changelog
|
||||
|
||||
### v3.3.1 (Proxy Routing Fix)
|
||||
|
||||
- **🌐 Proxy support**: the dashboard, gallery, tag page and metadata viewer now work when ComfyUI is reached through a proxy mount such as SwarmUI's `/ComfyBackendDirect/`. Scripts, images, API calls, progress streams and gallery folder filters stay under the mount, and the canvas node's buttons use ComfyUI's `fileURL()`. Direct ComfyUI URLs are unchanged. Thanks to @ghosttigerllc-bit ([#164](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/164))
|
||||
- **📦 Dependencies**: the runtime minimums stay at `Pillow>=8.0.0` and `watchdog>=2.1.0`, so updating Prompt Manager never forces upgrades in the shared ComfyUI environment
|
||||
|
||||
### v3.3.0 (Hardening, Background Jobs & Worker Threads)
|
||||
|
||||
- **🔒 Security**: gallery roots restricted to ComfyUI directories and validated wherever they are read; media-only, fail-closed image serving; Windows path escape in static routes closed; quote-safe HTML escaping with inline handlers moved to data attributes; the CivitAI key is never returned and only sent to HTTPS CivitAI hosts; auto-tag paths contained; side-effect routes are POST-only; backups include the WAL and restores are size-capped, integrity-checked and applied through SQLite's online backup API
|
||||
- **🧵 Worker Threads**: a Settings slider (default half the cores) used by the image rescan, thumbnail generation and duplicate detection. Measured: a 1 GB sample scanned in 2.5 s instead of 18 s at 8 threads
|
||||
- **🔁 Background Jobs**: the image rescan and the duplicate scan run on the server, stream progress, survive a closed tab, minimize to a progress pill and re-attach after a reload
|
||||
- **🎬 Video Prompts**: the scan reads the prompt ComfyUI's video savers embed in MP4/MOV/WebM containers (needs `ffprobe`)
|
||||
- **🖼️ Thumbnails**: written atomically and served under versioned URLs so a regenerated thumbnail never shows a stale cached copy; the per-run cap applies only to thumbnails that still need generating and the gallery says how many remain
|
||||
- **🏷️ Tag Autocomplete**: the dashboard tag box suggests tags, narrowed to those that co-occur with the tags already typed
|
||||
- **♾️ Infinite Scrolling**: optional in the gallery and the dashboard list (saved default plus a session toggle)
|
||||
- **⚙️ Gallery Settings** now apply at load; grid columns, media info and the Rescan button work as labelled
|
||||
- **🐛 Correctness**: run counts increment when ComfyUI dequeues a prompt; negative prompts are never traced; per-thread SQLite connections and idempotent migrations (an old rebuild could delete every image row); `truncate_logs` no longer rolls the log over; the ComfyUI canvas loads only the canvas extension; Windows console banners are ASCII-safe; subfolders use `/` on every platform
|
||||
- **🧪 Tests & CI**: 1,460 Python and 142 JavaScript tests (from 510 and 26), coverage 96% with an 85% gate, CI on Linux, Windows and macOS for Python 3.10–3.12, flake8 clean
|
||||
- **⚠️ API changes**: image paths in responses are relative to the ComfyUI tree (use `url`); `autotag/download`, `autotag/start` and the thumbnail progress stream are POST-only; `lora/status` returns `has_civitai_api_key` instead of the key; `/search` is paginated
|
||||
|
||||
### v3.2.4 (Recently Used Sorting & Queue-Time Run Counts)
|
||||
|
||||
- **🕒 Recently Used Sort**: sort the dashboard by last use, most used or rating, server-side across pages
|
||||
- **🔢 Run Counts**: `run_count` and `last_used_at` are recorded through the ComfyUI queue hook, so cached nodes still count
|
||||
- **🔗 Positive-Prompt Image Linking**: generated images link to the positive prompt traced through the queued graph
|
||||
|
||||
### v3.2.1 (LoRA Manager Integration)
|
||||
|
||||
- **🔗 LoRA Manager Integration**: Import LoRA metadata, trigger words, and CivitAI example images from [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manager) into your prompt database
|
||||
- **💉 Auto-Inject Trigger Words**: Optionally append LoRA trigger words when `<lora:name:weight>` is detected in prompts during encoding
|
||||
- **🔑 CivitAI API Key Support**: Authenticate with CivitAI to download NSFW preview images
|
||||
- **📥 Import Progress Modal**: Real-time SSE streaming progress during LoRA import with per-model status
|
||||
|
||||
> **Note:** LoRA Manager support is a WIP — this was a highly requested feature. Please [open issues](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues) for any bugs or feature requests.
|
||||
|
||||
### v3.2.0 (Folder Filter & QoL Improvements)
|
||||
|
||||
- **📁 Folder Filter**: New folder dropdown in the search panel to filter prompts by output subdirectory
|
||||
- **📂 Multi-Directory Gallery Scan**: Configure multiple output directories and browse them all in one gallery
|
||||
- **🖼️ Filmstrip Prompt Display**: Show prompt text and copy button in the filmstrip image viewer
|
||||
- **🖱️ Click-to-Close Viewer**: Click outside the image viewer to close it
|
||||
- **🐛 Bug Fixes**: Fixed database path config, UnboundLocalError in text inputs, node caching skip, diagnostics singleton lifecycle
|
||||
|
||||
### v3.1.0 (WD14 Tagger, Tailwind v4 & Major Refactors)
|
||||
|
||||
- **🏷️ WD14 Tagger Support**: Added WD14 SwinV2 and WD14 ViT as new auto-tag model options — fast ONNX-based classifiers (~400MB) that output Danbooru tags with confidence scores, no prompt needed
|
||||
- **🎚️ Threshold Controls**: Adjustable general tag (default 0.35) and character tag (default 0.85) confidence thresholds for WD14 models via UI sliders
|
||||
- **🎨 Tailwind v4 Migration**: Complete frontend redesign with ComfyUI theme token system (`--pm-*` design tokens) that adapts to ComfyUI's color palette
|
||||
- **🏗️ API Refactor**: Split monolithic `api.py` into domain-specific route modules for better maintainability
|
||||
- **🗃️ Junction Table Tags**: Migrated tag storage from JSON column to normalized junction tables with proper foreign keys
|
||||
- **📄 Tags Page**: Dedicated tag management page with search, rename, merge, and delete operations
|
||||
- **🔍 Load More Button**: Progressive prompt loading instead of traditional pagination
|
||||
- **🔒 Security Hardening**: Resolved 15 GitHub code scanning alerts, prevented stack trace exposure in SSE errors
|
||||
- **🧪 Test Coverage**: Added comprehensive unit tests for validators, hashing, metadata, config, and prompt tracker; API integration tests; database tests with CI pipeline
|
||||
- **⚙️ CI/CD Improvements**: Added code quality scanning, security scanning, Dependabot, and pre-commit hooks with Black formatting
|
||||
- **🐛 Bug Fixes**: Fixed custom image scan directory config (#108), gallery crash on thumbnail linking, broken imports after API refactor, widget value sync issues, and prepend/append metadata tag filtering in auto-tagger
|
||||
|
||||
### v3.0.26 (Top Bar Button)
|
||||
|
||||
- **🔘 Top Bar PM Button**: Added a "PM" button to the ComfyUI top bar for quick access to the admin dashboard
|
||||
- **🚀 No Node Required**: Open PromptManager without having to add the node to your workflow first
|
||||
- **🎨 Teal Button Design**: Distinctive teal-colored button positioned before the settings gear
|
||||
- **🔄 Frontend Compatibility**: Works with both new ComfyUI frontend (v1.3.1+) and legacy versions
|
||||
|
||||
### v3.0.25 (Film Strip & QoL Improvements)
|
||||
|
||||
- **🎞️ Film Strip Preview**: Backported from V2 - thumbnail strip under each prompt showing linked images
|
||||
- **🖼️ Quick Image Viewer**: Click any filmstrip thumbnail to open full-size image viewer
|
||||
- **📸 Smart Thumbnails**: Automatically uses generated thumbnails for faster loading, falls back to full images
|
||||
- **⬆️ Auto-Scroll on Pagination**: Page navigation now smoothly scrolls to top after loading
|
||||
|
||||
### v3.0.24 (Manual Prompt Creation)
|
||||
|
||||
- **➕ Add Prompt Button**: New "Add Prompt" button in the admin dashboard Bulk Actions bar
|
||||
- **📝 Direct Prompt Creation**: Create prompts directly from the UI without generating images
|
||||
- **🛡️ Protection from Pruning**: New prompts are automatically protected from orphan cleanup by default
|
||||
- **🏷️ Full Metadata Support**: Set category, tags, rating, and notes when creating prompts
|
||||
- **🔒 Orphan Protection Toggle**: Option to protect or unprotect prompts from maintenance cleanup
|
||||
|
||||
### v3.0.23 (AI AutoTag Release)
|
||||
|
||||
- **🏷️ AI AutoTag**: Automatically tag your entire image collection using JoyCaption vision models
|
||||
- **🤖 Dual Model Support**: Choose between JoyCaption Beta One FP16 (high quality) or GGUF FP8 (lower VRAM)
|
||||
- **📋 Two Tagging Modes**: Batch mode for bulk tagging, Review mode for one-by-one approval
|
||||
- **⚙️ Smart Tag Handling**: Skip already-tagged images or re-tag all - respects your existing work
|
||||
- **✏️ Customizable Prompts**: Adjust the system prompt to match your collection style and tag vocabulary
|
||||
- **🔄 Real-time Progress**: SSE-based streaming progress updates during batch operations
|
||||
- **📦 Tag Accordion UI**: Collapsible tag display shows first row with expandable section for large tag sets
|
||||
- **🎯 Review Confirmation**: Prompts before overwriting existing tags in review mode
|
||||
|
||||
### v3.0.10 (Batch Processing Node)
|
||||
|
||||
- **📋 PromptSearchList Node**: New node that searches the prompt database and outputs results as a list
|
||||
- **🔗 Batch Processing Support**: Uses `OUTPUT_IS_LIST` for native ComfyUI batch processing compatibility
|
||||
- **🔍 Full Search Capabilities**: Search by text, category, tags, minimum rating, and result limit
|
||||
- **⚡ Lightweight Operation**: Read-only search with no database writes
|
||||
- **🎯 Workflow Integration**: Connect directly to batch processing nodes like String OutputList
|
||||
|
||||
### v3.0.9 (Output-Wide Gallery & Video Support)
|
||||
|
||||
- **📁 Output-Wide Gallery Mode**: Complete gallery redesign to show ALL content from ComfyUI output directory
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
# 🔐 Security Policy
|
||||
|
||||
## Supported Versions
|
||||
|
||||
Security updates are provided for the actively maintained versions of
|
||||
**ComfyUI_PromptManager**.
|
||||
|
||||
| Version | Supported |
|
||||
|-----------------|-----------|
|
||||
| `main` / latest | ✅ |
|
||||
| Older releases | ❌ |
|
||||
|
||||
Users are encouraged to stay up to date with the latest version.
|
||||
|
||||
---
|
||||
|
||||
## Reporting a Vulnerability
|
||||
|
||||
For most security-related issues, **please open a GitHub issue** in this
|
||||
repository.
|
||||
|
||||
When opening an issue:
|
||||
- Clearly describe the problem
|
||||
- Include steps to reproduce
|
||||
- Mention the affected version or commit
|
||||
- Avoid posting secrets, tokens, or private data
|
||||
|
||||
### Sensitive Issues
|
||||
|
||||
If you believe the issue:
|
||||
- Could enable remote code execution
|
||||
- Exposes credentials or private data
|
||||
- Is actively exploitable in the wild
|
||||
|
||||
Please use **GitHub Security Advisories** instead:
|
||||
- Go to the **Security** tab
|
||||
- Click **Report a vulnerability**
|
||||
|
||||
---
|
||||
|
||||
## Response Expectations
|
||||
|
||||
- Issues are typically reviewed within **a few days**
|
||||
- Valid issues will be labeled and tracked publicly
|
||||
- Fixes may be discussed openly and merged transparently
|
||||
|
||||
---
|
||||
|
||||
## Scope
|
||||
|
||||
This policy applies to:
|
||||
- Code contained in this repository
|
||||
- Default configurations and documented usage
|
||||
|
||||
It does not cover:
|
||||
- Third-party nodes, models, or extensions
|
||||
- User workflows, prompts, or local environment issues
|
||||
|
||||
---
|
||||
|
||||
## Thank You
|
||||
|
||||
Community reports and transparency help keep the project healthy.
|
||||
Thanks for contributing responsibly.
|
||||
@@ -0,0 +1,130 @@
|
||||
# Testing and merge confidence
|
||||
|
||||
Routine changes should merge after the required checks pass and the diff is
|
||||
reviewed. Add a regression test for each bug; prove it fails before the fix.
|
||||
|
||||
## What runs
|
||||
|
||||
| Layer | Real components | Replaced boundaries |
|
||||
| --- | --- | --- |
|
||||
| Python tests | Domain logic, SQLite, API handlers, filesystem checks | ComfyUI host and ML runtimes where needed |
|
||||
| JavaScript tests | Frontend modules and controller methods | DOM/network fixtures for focused edge cases |
|
||||
| Viewer checks | Production Viewer.js and frontend methods in Chromium desktop/touch | Controlled images and page fixtures |
|
||||
| Real ComfyUI | Pinned CPU host, actual extension, node definitions, widgets, routes and SQLite | Inference is not queued; browser startup settings and host preference writes are isolated |
|
||||
| Page workflows | Production HTML, startup, click handlers, HTTP routes, SQLite, PNG files, thumbnails | ComfyUI host services and the tagging model |
|
||||
|
||||
Page workflows cover save/edit/reload, duplicate-edit rejection, gallery image
|
||||
identity, Tags selection and viewer cleanup, delayed responses, auto-tag review
|
||||
and persistence, and infinite scrolling without replacing existing cards.
|
||||
Data safety workflows download and upload actual SQLite backups, verify restored
|
||||
prompts/tags/image links, reject corrupt uploads without changing live data, and
|
||||
exercise cancellation and network-failure recovery for restores and deletion.
|
||||
Duplicate cleanup verifies only the selected file disappears; prompt deletion
|
||||
verifies the other prompt and original image files survive.
|
||||
They run in Chromium, Firefox, WebKit, and Chromium touch emulation. No GPU,
|
||||
model downloads, running ComfyUI instance, or personal configuration is required.
|
||||
The server binds to localhost and creates disposable directories and a database;
|
||||
its reset endpoint exists only in the test harness.
|
||||
|
||||
## Run locally
|
||||
|
||||
Use a Python environment containing `requirements-dev.txt` and Node.js 24:
|
||||
|
||||
```sh
|
||||
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
|
||||
npm run --prefix tools/vendor test:coverage
|
||||
npm run --prefix tools/vendor test:workflows -- --project firefox --project webkit
|
||||
```
|
||||
|
||||
`PM_TEST_PYTHON` selects the Python executable; `PM_TEST_PORT` selects the local
|
||||
port (default 8765). Set these environment variables when Python is outside PATH
|
||||
or the port is occupied. Test servers are never reused. On unsupported Linux
|
||||
distributions, browser system dependencies may need manual installation; CI uses
|
||||
Ubuntu. For a quick browser run, use `test:workflows -- --project chromium`.
|
||||
|
||||
## Managed real ComfyUI
|
||||
|
||||
The shared engine comes from an exact commit of
|
||||
[ComfyAssets Toolkit](https://github.com/ComfyAssets/ComfyAssets-Toolkit), recorded
|
||||
in `tools/vendor/package.json` and its integrity-checked lockfile. The local
|
||||
adapter sets PromptManager's database/configuration and monitoring policy;
|
||||
source discovery, coverage mapping/gates and process lifecycle are shared code.
|
||||
KikoTools remains the reference and is unchanged by this integration.
|
||||
|
||||
Use a **dedicated disposable** ComfyUI checkout at the revision in
|
||||
`tools/vendor/comfyui-runtime.json`. With a separate Python environment, install
|
||||
that checkout's requirements, this repo's `requirements-dev.txt`, and the exact
|
||||
CPU torch/torchvision pins from the JSON (using PyTorch's CPU wheel index).
|
||||
Install Chromium using the Playwright command above, then run:
|
||||
|
||||
```sh
|
||||
PM_COMFY_ROOT=/path/to/disposable/ComfyUI \
|
||||
PM_TEST_PYTHON=/path/to/cpu-env/bin/python \
|
||||
npm run --prefix tools/vendor test:comfyui
|
||||
```
|
||||
|
||||
CI performs this setup automatically. `PM_COMFY_URL` can select another
|
||||
`http://127.0.0.1:PORT/` address (default port 8189). This is an owned managed
|
||||
server: a running host is never reused. Do not point it at your live checkout.
|
||||
The launcher verifies the upstream SHA and clean tracked files, snapshots only
|
||||
tracked PromptManager files, and creates temporary user/input/output/database/
|
||||
configuration/log directories. Personal untracked files are excluded. It runs
|
||||
CPU/offline with only PromptManager enabled and removes owned state on shutdown.
|
||||
The upstream checkout must also be disposable: ComfyUI may initialize local
|
||||
runtime files there. Browser host-settings writes are intercepted; actual
|
||||
PromptManager routes use the disposable host and database.
|
||||
|
||||
`real-comfyui-reports` retains startup logs, version attachments, screenshots and
|
||||
traces for 14 days. Update upstream/frontend/CPU pins together after reviewing
|
||||
and running both startup projects; never replace the SHA with a moving branch.
|
||||
Update the toolkit via an exact reviewed commit archive and commit both npm
|
||||
manifest and lockfile. No toolkit or npm installation is needed by end users.
|
||||
|
||||
## Coverage and failures
|
||||
|
||||
Python CI enforces 85% combined statement/branch coverage. Frontend coverage merges
|
||||
Node V8 unit coverage with Chromium page-workflow coverage through c8. Every
|
||||
`web/js/**/*.js` and `web/js/**/*.mjs` file appears, including untouched files at zero coverage. The
|
||||
controller tests provide V8 filenames so their executed code is counted too.
|
||||
|
||||
`tools/vendor/frontend-coverage-baseline.json` holds a floor for every file and
|
||||
metric. CI rejects regressions, missing measurements, and new files without a
|
||||
reviewed baseline. These floors reflect current coverage, with two percentage
|
||||
points of allowance for runtime/timing differences; they do not claim 80% coverage
|
||||
of the large controllers. New files need at least 80% line coverage when added via
|
||||
the baseline updater. Raise floors with `npm run --prefix tools/vendor
|
||||
coverage:update`; it reruns the suites and cannot lower an existing floor.
|
||||
Review baseline changes like code changes.
|
||||
|
||||
The `viewerjs` check runs on every PR, including backend-only changes. GitHub
|
||||
artifacts named `frontend-test-reports` contain HTML coverage, a JSON summary,
|
||||
and failure screenshots/traces. Open traces with Playwright's `show-trace` tool.
|
||||
|
||||
The real-host suite covers both Classic and Nodes 2.0 startup, switching both
|
||||
ways, actual database prompt selection, execution-text serialization, workflow
|
||||
reload, widget/style cleanup, node schemas and standalone-page opening. It does
|
||||
not queue inference or certify every existing feature, arbitrary host versions,
|
||||
GPU/model changes, cross-pack installation or visual design changes. Those still
|
||||
need targeted validation.
|
||||
|
||||
## Required merge checks
|
||||
|
||||
`.github/branch-protection.json` records the required checks: all twelve Python/OS
|
||||
combinations, Python coverage, formatting, frontend modules, `viewerjs`, and `real ComfyUI compatibility`.
|
||||
The branch must be current with main, and administrators follow the same checks.
|
||||
It does not require an additional human approval.
|
||||
|
||||
A repository administrator applies the configuration after the new workflow is
|
||||
passing, then verifies it:
|
||||
|
||||
```sh
|
||||
gh api --method PUT repos/ComfyAssets/ComfyUI_PromptManager/branches/main/protection \
|
||||
--input .github/branch-protection.json
|
||||
gh api repos/ComfyAssets/ComfyUI_PromptManager/branches/main/protection/required_status_checks
|
||||
```
|
||||
|
||||
Update the manifest whenever required job names change. Keep the frontend check
|
||||
free of path filters so a required check is never silently skipped.
|
||||
@@ -1,19 +1,21 @@
|
||||
"""
|
||||
ComfyUI_PromptManager: A comprehensive ComfyUI custom node that extends text encoding
|
||||
with persistent prompt storage, advanced search capabilities, and an automatic image
|
||||
ComfyUI_PromptManager: A comprehensive ComfyUI custom node that extends text encoding
|
||||
with persistent prompt storage, advanced search capabilities, and an automatic image
|
||||
gallery system.
|
||||
|
||||
This module provides two main node types:
|
||||
This module provides prompt node types:
|
||||
- PromptManager: CLIP encoding node that outputs CONDITIONING
|
||||
- PromptManagerText: Text-only node that outputs STRING with prepend/append functionality
|
||||
- PromptManagerFocused: Editor-first layout using the same encoding contract
|
||||
- PromptManagerText: Text-only node that outputs STRING with prepend/append
|
||||
functionality
|
||||
|
||||
Both nodes share the same database backend for persistent prompt storage and include
|
||||
The nodes share the same database backend for persistent prompt storage and include
|
||||
an automatic image gallery system that monitors ComfyUI output directories and links
|
||||
generated images to their source prompts.
|
||||
|
||||
Features:
|
||||
- SQLite-based persistent prompt storage with deduplication
|
||||
- Advanced search and categorization system
|
||||
- Advanced search and categorization system
|
||||
- Real-time image gallery with metadata extraction
|
||||
- Web-based admin dashboard for prompt management
|
||||
- Comprehensive logging and diagnostics
|
||||
@@ -37,21 +39,28 @@ def get_version():
|
||||
return "unknown"
|
||||
|
||||
|
||||
from .prompt_manager import PromptManager
|
||||
from .prompt_manager_text import PromptManagerText
|
||||
from .prompt_manager import PromptManager, PromptManagerFocused # noqa: E402
|
||||
from .prompt_manager_text import PromptManagerText # noqa: E402
|
||||
from .prompt_search_list import PromptSearchList # noqa: E402
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PromptManager": PromptManager,
|
||||
"PromptManagerFocused": PromptManagerFocused,
|
||||
"PromptManagerText": PromptManagerText,
|
||||
"PromptSearchList": PromptSearchList,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PromptManager": "Prompt Manager",
|
||||
"PromptManager": "Prompt Manager (Classic)",
|
||||
"PromptManagerFocused": "Prompt Manager (Focused)",
|
||||
"PromptManagerText": "Prompt Manager Text",
|
||||
"PromptSearchList": "Prompt Search List",
|
||||
}
|
||||
|
||||
# Define path to web directory for UI components
|
||||
WEB_DIRECTORY = "web"
|
||||
# Canvas extension only. ComfyUI loads every .js under this directory into the
|
||||
# canvas page, so the admin/gallery bundles and vendored libraries under web/
|
||||
# are served by our own routes instead (see py/api).
|
||||
WEB_DIRECTORY = "web/comfy"
|
||||
|
||||
# Add API routes (same pattern as ComfyUI_Assets)
|
||||
try:
|
||||
@@ -76,15 +85,77 @@ except Exception as e:
|
||||
|
||||
logger = get_logger("prompt_manager.init")
|
||||
logger.error(f"Failed to register API routes: {e}")
|
||||
except:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Shared database, set up with the monitoring system below. Stays None when that
|
||||
# fails so the queue hook can report the database as unavailable instead of crashing.
|
||||
_global_db = None
|
||||
|
||||
# Start image monitoring globally at module import time
|
||||
# This ensures images are linked regardless of which node is used in the workflow
|
||||
try:
|
||||
from .database.operations import PromptDatabase
|
||||
from .utils.image_monitor import get_image_monitor
|
||||
from .utils.prompt_tracker import get_prompt_tracker
|
||||
from .utils.logging_config import get_logger
|
||||
|
||||
_init_logger = get_logger("prompt_manager.init")
|
||||
_init_logger.info("Starting global image monitoring system...")
|
||||
|
||||
# Initialize database and get singleton instances
|
||||
_global_db = PromptDatabase()
|
||||
_global_prompt_tracker = get_prompt_tracker(_global_db)
|
||||
_global_image_monitor = get_image_monitor(_global_db, _global_prompt_tracker)
|
||||
|
||||
# Start monitoring
|
||||
_global_image_monitor.start_monitoring()
|
||||
_init_logger.info("Global image monitoring system started successfully")
|
||||
|
||||
except Exception as e:
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
|
||||
_init_logger = get_logger("prompt_manager.init")
|
||||
_init_logger.error(f"Failed to start image monitoring: {e}")
|
||||
except Exception:
|
||||
print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}")
|
||||
|
||||
|
||||
def _db_factory():
|
||||
"""The shared database for the queue hook; called lazily per request."""
|
||||
if _global_db is None:
|
||||
raise RuntimeError("PromptManager database unavailable")
|
||||
return _global_db
|
||||
|
||||
|
||||
# Count prompt runs when ComfyUI dequeues workflows (it skips unchanged nodes, so
|
||||
# node execution alone misses re-runs). Separate from monitoring so either can fail.
|
||||
try:
|
||||
from .py.config import server_instance
|
||||
from .utils.usage_tracking import register_queue_hook
|
||||
|
||||
register_queue_hook(server_instance, _db_factory)
|
||||
except Exception as e:
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
|
||||
get_logger("prompt_manager.init").error(f"Failed to register usage hook: {e}")
|
||||
except Exception:
|
||||
import logging
|
||||
|
||||
logging.getLogger("prompt_manager.init").error(
|
||||
f"Failed to register usage hook: {e}"
|
||||
)
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
# Print startup message with loaded tools
|
||||
print()
|
||||
print(f"\033[94m[ComfyUI-PromptManager] Version:\033[0m {get_version()}")
|
||||
for node_key, display_name in NODE_DISPLAY_NAME_MAPPINGS.items():
|
||||
print(f"🫶 \033[94mLoaded:\033[0m {display_name}")
|
||||
# ASCII only: a Windows console using a legacy code page cannot print emoji,
|
||||
# and a UnicodeEncodeError here would abort loading the whole node.
|
||||
print(f" \033[94mLoaded:\033[0m {display_name}")
|
||||
print(f"\033[94mTotal: {len(NODE_CLASS_MAPPINGS)} tools loaded\033[0m")
|
||||
print()
|
||||
|
||||
@@ -8,4 +8,4 @@ search capabilities, metadata management, and image tracking functionality.
|
||||
from .operations import PromptDatabase
|
||||
from .models import PromptModel
|
||||
|
||||
__all__ = ["PromptDatabase", "PromptModel"]
|
||||
__all__ = ["PromptDatabase", "PromptModel"]
|
||||
|
||||
@@ -2,62 +2,167 @@
|
||||
Database schema and models for KikoTextEncode prompt storage.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import datetime
|
||||
import shutil
|
||||
import sqlite3
|
||||
import os
|
||||
from typing import Optional
|
||||
import threading
|
||||
import weakref
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from typing import Set, Tuple
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, current_dir)
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
# last_used_at is UTC ISO text so ORDER BY can compare it as text. The backfill
|
||||
# writes milliseconds (SQLite's %f); live writes use microseconds so runs in the
|
||||
# same millisecond still order correctly. Both forms sort correctly together.
|
||||
ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
|
||||
|
||||
|
||||
def utc_now_iso() -> str:
|
||||
"""Current UTC time as ISO text with microseconds.
|
||||
|
||||
Example: 2026-10-07T18:39:06.758123+00:00
|
||||
"""
|
||||
return datetime.datetime.now(datetime.timezone.utc).isoformat(
|
||||
timespec="microseconds"
|
||||
)
|
||||
|
||||
|
||||
def normalize_image_path(path: str) -> str:
|
||||
"""Canonical form of an image path for generated_images.file_path.
|
||||
|
||||
Absolute, with '..' and '.' collapsed and the platform's case folding
|
||||
applied, so two spellings of one file (or the same basename in two date
|
||||
folders) compare the way the filesystem would.
|
||||
"""
|
||||
return os.path.normcase(os.path.normpath(os.path.abspath(path)))
|
||||
|
||||
|
||||
# How long a connection waits for a locked database before giving up. Thirty
|
||||
# seconds covers a scan or a thumbnail run writing while ComfyUI saves images,
|
||||
# and the slow disks of CI runners.
|
||||
BUSY_TIMEOUT_SECONDS = 30
|
||||
SQLITE_HEADER = b"SQLite format 3\x00"
|
||||
REQUIRED_PROMPT_COLUMNS = ("id", "text", "created_at")
|
||||
|
||||
|
||||
class PromptModel:
|
||||
"""Database model for prompt storage and schema management."""
|
||||
|
||||
def __init__(self, db_path: str = "prompts.db"):
|
||||
|
||||
# Paths (normalised, see _schema_key) whose schema has been created and
|
||||
# migrated in this process. Many PromptDatabase instances point at the
|
||||
# same file; only the first one pays for the schema check.
|
||||
_initialized_paths: Set[str] = set()
|
||||
_init_lock = threading.Lock()
|
||||
# Every live model, so a shutdown or a test teardown can close them all
|
||||
# before the database file is deleted (Windows refuses otherwise).
|
||||
_instances: "weakref.WeakSet[PromptModel]" = weakref.WeakSet()
|
||||
_instances_lock = threading.Lock()
|
||||
|
||||
def __init__(self, db_path: str):
|
||||
"""
|
||||
Initialize the database model.
|
||||
|
||||
|
||||
Args:
|
||||
db_path: Path to the SQLite database file
|
||||
db_path: Absolute path to the SQLite database file
|
||||
"""
|
||||
self.logger = get_logger('prompt_manager.database.models')
|
||||
self.logger = get_logger("prompt_manager.database.models")
|
||||
self.logger.debug(f"Initializing database model with path: {db_path}")
|
||||
self.db_path = db_path
|
||||
# One connection per thread (see get_connection); every open connection
|
||||
# is also tracked here so restore_from_file can close them all.
|
||||
self._local = threading.local()
|
||||
self._connections: Set[sqlite3.Connection] = set()
|
||||
self._conn_lock = threading.Lock()
|
||||
# Bumped by _close_all_connections so threads holding a closed
|
||||
# connection reopen on their next get_connection().
|
||||
self._generation = 0
|
||||
with PromptModel._instances_lock:
|
||||
PromptModel._instances.add(self)
|
||||
self._ensure_database_exists()
|
||||
|
||||
|
||||
def _ensure_database_exists(self) -> None:
|
||||
"""
|
||||
Create database and tables if they don't exist.
|
||||
|
||||
|
||||
Sets up the database schema including tables for prompts and generated images,
|
||||
creates necessary indexes, and applies any pending migrations.
|
||||
|
||||
|
||||
Raises:
|
||||
Exception: If database creation fails
|
||||
"""
|
||||
key = self._schema_key()
|
||||
with PromptModel._init_lock:
|
||||
if key in PromptModel._initialized_paths and os.path.exists(self.db_path):
|
||||
return
|
||||
try:
|
||||
conn = sqlite3.connect(self.db_path, timeout=BUSY_TIMEOUT_SECONDS)
|
||||
try:
|
||||
conn.execute("PRAGMA journal_mode = WAL")
|
||||
conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
self._create_tables(conn)
|
||||
self._create_indexes(conn)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error creating database: {e}")
|
||||
raise
|
||||
PromptModel._initialized_paths.add(key)
|
||||
|
||||
def _schema_key(self) -> str:
|
||||
"""Normalised identity of db_path for the once-per-process schema cache."""
|
||||
return os.path.normcase(os.path.realpath(self.db_path))
|
||||
|
||||
@classmethod
|
||||
def reset_schema_cache(cls) -> None:
|
||||
"""Forget which databases were initialised (tests and restores)."""
|
||||
with cls._init_lock:
|
||||
cls._initialized_paths.clear()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _table_rebuild(self, conn: sqlite3.Connection):
|
||||
"""Run a copy-and-swap table rebuild with foreign keys disabled.
|
||||
|
||||
DROP TABLE on a parent table with foreign_keys ON performs an implicit
|
||||
DELETE that cascades into child rows (generated_images, prompt_tags),
|
||||
so the rebuild runs with the pragma off. The pragma is a no-op inside
|
||||
a transaction, hence the commits around it. Errors roll back and
|
||||
propagate after the pragma is restored.
|
||||
"""
|
||||
conn.commit()
|
||||
conn.execute("PRAGMA foreign_keys = OFF")
|
||||
# Explicit BEGIN: in sqlite3's legacy transaction mode DDL autocommits,
|
||||
# so without it a failed copy would leave the *_new table behind.
|
||||
conn.execute("BEGIN")
|
||||
try:
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
self._create_tables(conn)
|
||||
self._create_indexes(conn)
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error creating database: {e}")
|
||||
yield
|
||||
conn.commit()
|
||||
except BaseException:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
finally:
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
|
||||
def _create_tables(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Create the prompts and generated_images tables with all required columns.
|
||||
|
||||
|
||||
Args:
|
||||
conn: Active database connection
|
||||
|
||||
|
||||
Creates:
|
||||
- prompts table: Stores prompt text and metadata
|
||||
- generated_images table: Links generated images to their source prompts
|
||||
@@ -72,10 +177,12 @@ class PromptModel:
|
||||
tags TEXT,
|
||||
rating INTEGER CHECK(rating >= 1 AND rating <= 5),
|
||||
notes TEXT,
|
||||
hash TEXT UNIQUE
|
||||
hash TEXT UNIQUE,
|
||||
last_used_at TIMESTAMP,
|
||||
run_count INTEGER NOT NULL DEFAULT 0
|
||||
)
|
||||
""")
|
||||
|
||||
|
||||
# Create images table for gallery functionality
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS generated_images (
|
||||
@@ -91,23 +198,51 @@ class PromptModel:
|
||||
workflow_data TEXT,
|
||||
prompt_metadata TEXT,
|
||||
parameters TEXT,
|
||||
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE
|
||||
file_path TEXT,
|
||||
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
|
||||
UNIQUE(prompt_id, file_path)
|
||||
)
|
||||
""")
|
||||
|
||||
|
||||
# Create normalized tag tables (junction table pattern)
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS tags (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL UNIQUE
|
||||
)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS prompt_tags (
|
||||
prompt_id INTEGER NOT NULL,
|
||||
tag_id INTEGER NOT NULL,
|
||||
PRIMARY KEY (prompt_id, tag_id),
|
||||
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (tag_id) REFERENCES tags(id) ON DELETE CASCADE
|
||||
)
|
||||
""")
|
||||
|
||||
# Check if we need to migrate from old schema with workflow_name
|
||||
self._migrate_workflow_name_removal(conn)
|
||||
|
||||
|
||||
# Fix foreign key data type mismatch
|
||||
self._migrate_foreign_key_types(conn)
|
||||
|
||||
|
||||
# Migrate JSON tags to normalized junction tables
|
||||
self._migrate_json_tags_to_junction(conn)
|
||||
|
||||
# Usage tracking columns (after any migration that rebuilds prompts)
|
||||
self._migrate_add_usage_columns(conn)
|
||||
|
||||
# Image uniqueness by full path instead of basename (3.2.4)
|
||||
self._migrate_image_uniqueness_by_path(conn)
|
||||
|
||||
def _create_indexes(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Create indexes for better query performance.
|
||||
|
||||
|
||||
Args:
|
||||
conn: Active database connection
|
||||
|
||||
|
||||
Creates indexes on:
|
||||
- Text content for search operations
|
||||
- Categories and tags for filtering
|
||||
@@ -115,50 +250,124 @@ class PromptModel:
|
||||
- Hash values for duplicate detection
|
||||
"""
|
||||
indexes = [
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)",
|
||||
("idx_prompts_text", "prompts(text)"),
|
||||
("idx_prompts_category", "prompts(category)"),
|
||||
("idx_prompts_created_at", "prompts(created_at)"),
|
||||
("idx_prompts_hash", "prompts(hash)"),
|
||||
("idx_prompts_rating", "prompts(rating)"),
|
||||
("idx_prompts_last_used", "prompts(last_used_at)"),
|
||||
("idx_prompts_run_count", "prompts(run_count)"),
|
||||
("idx_prompt_images", "generated_images(prompt_id)"),
|
||||
("idx_image_path", "generated_images(image_path)"),
|
||||
("idx_image_file_path", "generated_images(file_path)"),
|
||||
("idx_generation_time", "generated_images(generation_time)"),
|
||||
("idx_prompt_tags_tag", "prompt_tags(tag_id)"),
|
||||
("idx_tags_name", "tags(name)"),
|
||||
]
|
||||
|
||||
for index_sql in indexes:
|
||||
conn.execute(index_sql)
|
||||
|
||||
|
||||
for name, target in indexes:
|
||||
conn.execute(f"CREATE INDEX IF NOT EXISTS {name} ON {target}")
|
||||
|
||||
def get_connection(self) -> sqlite3.Connection:
|
||||
"""
|
||||
Get a database connection with proper configuration.
|
||||
|
||||
Get the calling thread's database connection.
|
||||
|
||||
Each thread gets its own ``sqlite3.Connection`` (kept in
|
||||
``threading.local``), created on first use and reused by that thread
|
||||
thereafter. Every connection is configured with WAL journaling,
|
||||
``row_factory = sqlite3.Row``, ``PRAGMA foreign_keys = ON`` and
|
||||
``PRAGMA busy_timeout`` of BUSY_TIMEOUT_SECONDS so concurrent writers
|
||||
wait instead of failing. Callers use it as a context manager
|
||||
(``with conn:``), which commits or rolls back that thread's transaction
|
||||
and leaves the connection open. No lock is held while a connection is in use.
|
||||
|
||||
Returns:
|
||||
sqlite3.Connection: Configured database connection
|
||||
sqlite3.Connection: The current thread's configured connection
|
||||
"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.row_factory = sqlite3.Row # Enable dict-like access to rows
|
||||
conn = getattr(self._local, "conn", None)
|
||||
if conn is not None and self._local.generation == self._generation:
|
||||
return conn
|
||||
|
||||
# check_same_thread=False only so restore_from_file can close other
|
||||
# threads' connections; each connection is still used by one thread.
|
||||
conn = sqlite3.connect(
|
||||
self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
|
||||
)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode = WAL")
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
|
||||
self._local.conn = conn
|
||||
self._local.generation = self._generation
|
||||
with self._conn_lock:
|
||||
self._connections.add(conn)
|
||||
return conn
|
||||
|
||||
|
||||
@classmethod
|
||||
def close_all_instances(cls) -> None:
|
||||
"""Close every connection of every live model in this process."""
|
||||
with cls._instances_lock:
|
||||
models = list(cls._instances)
|
||||
for model in models:
|
||||
model.close_all()
|
||||
|
||||
def close_all(self) -> None:
|
||||
"""Close every thread's connection; they reopen lazily on next use.
|
||||
|
||||
Needed before deleting the database file (Windows keeps open files
|
||||
locked) and by tests that ran queries through executor threads.
|
||||
"""
|
||||
with self._conn_lock:
|
||||
connections = list(self._connections)
|
||||
self._connections.clear()
|
||||
self._generation += 1
|
||||
self._local.conn = None
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's connection, if it has one.
|
||||
|
||||
Other threads keep their connections. Call this before deleting or
|
||||
replacing the database file: Windows refuses to delete an open file.
|
||||
"""
|
||||
conn = getattr(self._local, "conn", None)
|
||||
if conn is None:
|
||||
return
|
||||
self._local.conn = None
|
||||
with self._conn_lock:
|
||||
self._connections.discard(conn)
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Remove workflow_name column if it exists in existing database.
|
||||
|
||||
|
||||
Args:
|
||||
conn: Active database connection
|
||||
|
||||
|
||||
This migration handles legacy schema updates by removing the deprecated
|
||||
workflow_name column while preserving all other data.
|
||||
|
||||
Raises:
|
||||
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
|
||||
"""
|
||||
try:
|
||||
# Check if workflow_name column exists
|
||||
cursor = conn.execute("PRAGMA table_info(prompts)")
|
||||
columns = [column[1] for column in cursor.fetchall()]
|
||||
|
||||
if 'workflow_name' in columns:
|
||||
self.logger.info("Migrating database: removing workflow_name column")
|
||||
|
||||
# Create new table without workflow_name
|
||||
if "workflow_name" not in columns:
|
||||
return
|
||||
|
||||
self.logger.info("Migrating database: removing workflow_name column")
|
||||
with self._table_rebuild(conn):
|
||||
# A previous attempt may have died between CREATE and RENAME
|
||||
conn.execute("DROP TABLE IF EXISTS prompts_new")
|
||||
conn.execute("""
|
||||
CREATE TABLE prompts_new (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
@@ -172,43 +381,43 @@ class PromptModel:
|
||||
hash TEXT UNIQUE
|
||||
)
|
||||
""")
|
||||
|
||||
# Copy data from old table to new table
|
||||
conn.execute("""
|
||||
INSERT INTO prompts_new (id, text, created_at, updated_at, category, tags, rating, notes, hash)
|
||||
SELECT id, text, created_at, updated_at, category, tags, rating, notes, hash
|
||||
INSERT INTO prompts_new (id, text, created_at, updated_at,
|
||||
category, tags, rating, notes, hash)
|
||||
SELECT id, text, created_at, updated_at, category, tags, rating,
|
||||
notes, hash
|
||||
FROM prompts
|
||||
""")
|
||||
|
||||
# Drop old table and rename new one
|
||||
conn.execute("DROP TABLE prompts")
|
||||
conn.execute("ALTER TABLE prompts_new RENAME TO prompts")
|
||||
|
||||
self.logger.info("Database migration completed")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Migration error: {e}")
|
||||
# If migration fails, the table creation will handle it
|
||||
|
||||
self.logger.info("Database migration completed")
|
||||
|
||||
except sqlite3.Error as e:
|
||||
self.logger.error(f"Migration error (workflow_name removal): {e}")
|
||||
raise
|
||||
|
||||
def _migrate_foreign_key_types(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Fix foreign key data type mismatch in generated_images table.
|
||||
|
||||
|
||||
Args:
|
||||
conn: Active database connection
|
||||
|
||||
|
||||
Converts prompt_id from TEXT to INTEGER type to match the prompts table's
|
||||
primary key type, ensuring referential integrity.
|
||||
|
||||
Raises:
|
||||
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
|
||||
"""
|
||||
try:
|
||||
# Check if generated_images table exists and has TEXT prompt_id
|
||||
cursor = conn.execute("PRAGMA table_info(generated_images)")
|
||||
columns = {column[1]: column[2] for column in cursor.fetchall()}
|
||||
|
||||
if 'prompt_id' in columns and columns['prompt_id'] == 'TEXT':
|
||||
self.logger.info("Migrating foreign key types: prompt_id TEXT -> INTEGER")
|
||||
|
||||
# Create new table with correct types
|
||||
if columns.get("prompt_id") != "TEXT":
|
||||
return
|
||||
|
||||
self.logger.info("Migrating foreign key types: prompt_id TEXT -> INTEGER")
|
||||
with self._table_rebuild(conn):
|
||||
conn.execute("DROP TABLE IF EXISTS generated_images_new")
|
||||
conn.execute("""
|
||||
CREATE TABLE generated_images_new (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
@@ -226,103 +435,415 @@ class PromptModel:
|
||||
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE
|
||||
)
|
||||
""")
|
||||
|
||||
# Copy data, converting prompt_id from TEXT to INTEGER
|
||||
conn.execute("""
|
||||
INSERT INTO generated_images_new
|
||||
(id, prompt_id, image_path, filename, generation_time, file_size,
|
||||
INSERT INTO generated_images_new
|
||||
(id, prompt_id, image_path, filename, generation_time, file_size,
|
||||
width, height, format, workflow_data, prompt_metadata, parameters)
|
||||
SELECT id, CAST(prompt_id AS INTEGER), image_path, filename, generation_time,
|
||||
file_size, width, height, format, workflow_data, prompt_metadata, parameters
|
||||
SELECT id, CAST(prompt_id AS INTEGER), image_path, filename,
|
||||
generation_time,
|
||||
file_size, width, height, format, workflow_data,
|
||||
prompt_metadata, parameters
|
||||
FROM generated_images
|
||||
WHERE prompt_id != '' AND prompt_id IS NOT NULL
|
||||
AND CAST(prompt_id AS INTEGER) IN (SELECT id FROM prompts)
|
||||
""")
|
||||
|
||||
# Drop old table and rename new one
|
||||
conn.execute("DROP TABLE generated_images")
|
||||
conn.execute("ALTER TABLE generated_images_new RENAME TO generated_images")
|
||||
|
||||
self.logger.info("Foreign key migration completed")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Foreign key migration error: {e}")
|
||||
# If migration fails, continue with existing schema
|
||||
|
||||
def migrate_database(self) -> None:
|
||||
conn.execute(
|
||||
"ALTER TABLE generated_images_new RENAME TO generated_images"
|
||||
)
|
||||
self.logger.info("Foreign key migration completed")
|
||||
|
||||
except sqlite3.Error as e:
|
||||
self.logger.error(f"Migration error (foreign key types): {e}")
|
||||
raise
|
||||
|
||||
def _migrate_json_tags_to_junction(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Apply any pending database migrations.
|
||||
|
||||
This method serves as an entry point for future schema migrations.
|
||||
Add new migration logic here as the database evolves.
|
||||
Populate tags and prompt_tags tables from legacy JSON tags column.
|
||||
|
||||
Runs once: skips if the tags table already has data. Uses json_each()
|
||||
to extract tag names from the JSON arrays stored in prompts.tags.
|
||||
|
||||
Raises:
|
||||
sqlite3.Error: If the migration fails (logged)
|
||||
"""
|
||||
# Future migrations can be added here
|
||||
pass
|
||||
|
||||
try:
|
||||
cursor = conn.execute("SELECT COUNT(*) FROM tags")
|
||||
if cursor.fetchone()[0] > 0:
|
||||
return # Already migrated
|
||||
|
||||
cursor = conn.execute(
|
||||
"SELECT COUNT(*) FROM prompts "
|
||||
"WHERE tags IS NOT NULL AND tags != '' AND tags != '[]'"
|
||||
)
|
||||
if cursor.fetchone()[0] == 0:
|
||||
return # No tags to migrate
|
||||
|
||||
self.logger.info("Migrating JSON tags to normalized junction tables")
|
||||
|
||||
# Insert all unique tag names
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO tags (name) "
|
||||
"SELECT DISTINCT je.value FROM prompts, json_each(prompts.tags) AS je "
|
||||
"WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND "
|
||||
"prompts.tags != '[]'"
|
||||
)
|
||||
|
||||
# Populate junction table
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
|
||||
"SELECT p.id, t.id "
|
||||
"FROM prompts p, json_each(p.tags) AS je "
|
||||
"JOIN tags t ON t.name = je.value "
|
||||
"WHERE p.tags IS NOT NULL AND p.tags != '' AND p.tags != '[]'"
|
||||
)
|
||||
|
||||
tag_count = conn.execute("SELECT COUNT(*) FROM tags").fetchone()[0]
|
||||
link_count = conn.execute("SELECT COUNT(*) FROM prompt_tags").fetchone()[0]
|
||||
self.logger.info(
|
||||
f"Tag migration complete: {tag_count} unique tags, {link_count} "
|
||||
"prompt-tag links"
|
||||
)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
self.logger.error(f"Migration error (tag junction): {e}")
|
||||
raise
|
||||
|
||||
def _migrate_add_usage_columns(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Add last_used_at and run_count (3.2.4) and estimate usage for existing rows.
|
||||
|
||||
The backfill runs only when the columns are added, i.e. on the upgrade
|
||||
from a version without usage tracking. Afterwards a NULL last_used_at
|
||||
means "never run": saving a prompt does not stamp it, only
|
||||
record_prompt_use does, and "Recently Used" relies on that.
|
||||
|
||||
The estimate comes from linked images: run_count is the image count (at
|
||||
least 1) and last_used_at the newest image time, falling back to created_at.
|
||||
Timestamps are compared with julianday() because prompts use ISO 'T'
|
||||
timestamps while generated_images uses SQLite's CURRENT_TIMESTAMP format.
|
||||
"""
|
||||
try:
|
||||
columns = {row[1] for row in conn.execute("PRAGMA table_info(prompts)")}
|
||||
added = False
|
||||
if "last_used_at" not in columns:
|
||||
self.logger.info("Adding prompt usage column last_used_at")
|
||||
conn.execute("ALTER TABLE prompts ADD COLUMN last_used_at TIMESTAMP")
|
||||
added = True
|
||||
if "run_count" not in columns:
|
||||
self.logger.info("Adding prompt usage column run_count")
|
||||
conn.execute(
|
||||
"ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL "
|
||||
"DEFAULT 0"
|
||||
)
|
||||
added = True
|
||||
if not added:
|
||||
return
|
||||
|
||||
cursor = conn.execute(f"""
|
||||
UPDATE prompts SET
|
||||
run_count = MAX(run_count, 1, (
|
||||
SELECT COUNT(*) FROM generated_images gi
|
||||
WHERE gi.prompt_id = prompts.id
|
||||
)),
|
||||
last_used_at = strftime('{ISO_TIMESTAMP_FORMAT}', MAX(
|
||||
COALESCE(julianday(created_at), julianday('now')),
|
||||
COALESCE((
|
||||
SELECT MAX(julianday(gi.generation_time))
|
||||
FROM generated_images gi WHERE gi.prompt_id = prompts.id
|
||||
), 0)
|
||||
))
|
||||
WHERE last_used_at IS NULL
|
||||
""")
|
||||
if cursor.rowcount:
|
||||
self.logger.info(f"Estimated usage for {cursor.rowcount} prompts")
|
||||
except sqlite3.Error as e:
|
||||
self.logger.error(f"Migration error (usage columns): {e}")
|
||||
raise
|
||||
|
||||
def _migrate_image_uniqueness_by_path(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Rebuild generated_images so an image is unique per (prompt_id, file_path).
|
||||
|
||||
Older schemas were unique per (prompt_id, filename), i.e. the basename,
|
||||
so ComfyUI's same-named files in different date folders were silently
|
||||
dropped. file_path holds normalize_image_path(image_path); rows whose
|
||||
paths normalise to the same file collapse to the newest one.
|
||||
|
||||
Raises:
|
||||
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
|
||||
"""
|
||||
try:
|
||||
columns = {
|
||||
row[1] for row in conn.execute("PRAGMA table_info(generated_images)")
|
||||
}
|
||||
if "file_path" in columns:
|
||||
return
|
||||
|
||||
self.logger.info("Migrating database: image uniqueness by full path")
|
||||
with self._table_rebuild(conn):
|
||||
rows = conn.execute(
|
||||
"SELECT id, prompt_id, image_path FROM generated_images ORDER BY id"
|
||||
).fetchall()
|
||||
keep = {}
|
||||
for row_id, prompt_id, image_path in rows:
|
||||
keep[(prompt_id, normalize_image_path(image_path))] = row_id
|
||||
|
||||
conn.execute("DROP TABLE IF EXISTS generated_images_new")
|
||||
conn.execute("""
|
||||
CREATE TABLE generated_images_new (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
prompt_id INTEGER NOT NULL,
|
||||
image_path TEXT NOT NULL,
|
||||
filename TEXT NOT NULL,
|
||||
generation_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
file_size INTEGER,
|
||||
width INTEGER,
|
||||
height INTEGER,
|
||||
format TEXT,
|
||||
workflow_data TEXT,
|
||||
prompt_metadata TEXT,
|
||||
parameters TEXT,
|
||||
file_path TEXT,
|
||||
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE
|
||||
CASCADE,
|
||||
UNIQUE(prompt_id, file_path)
|
||||
)
|
||||
""")
|
||||
conn.executemany(
|
||||
"""
|
||||
INSERT INTO generated_images_new
|
||||
(id, prompt_id, image_path, filename, generation_time, file_size,
|
||||
width, height, format, workflow_data, prompt_metadata, parameters,
|
||||
file_path)
|
||||
SELECT id, prompt_id, image_path, filename, generation_time,
|
||||
file_size,
|
||||
width, height, format, workflow_data, prompt_metadata,
|
||||
parameters,
|
||||
?
|
||||
FROM generated_images WHERE id = ?
|
||||
""",
|
||||
[(file_path, row_id) for (_, file_path), row_id in keep.items()],
|
||||
)
|
||||
conn.execute("DROP TABLE generated_images")
|
||||
conn.execute(
|
||||
"ALTER TABLE generated_images_new RENAME TO generated_images"
|
||||
)
|
||||
dropped = len(rows) - len(keep)
|
||||
if dropped:
|
||||
self.logger.info(f"Collapsed {dropped} duplicate image rows")
|
||||
self.logger.info("Image path uniqueness migration completed")
|
||||
|
||||
except sqlite3.Error as e:
|
||||
self.logger.error(f"Migration error (image path uniqueness): {e}")
|
||||
raise
|
||||
|
||||
def vacuum_database(self) -> None:
|
||||
"""
|
||||
Optimize database by running VACUUM command.
|
||||
|
||||
|
||||
Reclaims unused space and defragments the database file,
|
||||
improving query performance and reducing file size.
|
||||
"""
|
||||
try:
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with closing(sqlite3.connect(self.db_path)) as conn:
|
||||
conn.execute("VACUUM")
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error vacuuming database: {e}")
|
||||
|
||||
|
||||
def get_database_info(self) -> dict:
|
||||
"""
|
||||
Get information about the database.
|
||||
|
||||
|
||||
Returns:
|
||||
dict: Database statistics and information
|
||||
"""
|
||||
try:
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.execute("SELECT COUNT(*) as total_prompts FROM prompts")
|
||||
total_prompts = cursor.fetchone()['total_prompts']
|
||||
|
||||
total_prompts = cursor.fetchone()["total_prompts"]
|
||||
|
||||
cursor = conn.execute(
|
||||
"SELECT COUNT(DISTINCT category) as unique_categories FROM prompts WHERE category IS NOT NULL"
|
||||
"SELECT COUNT(DISTINCT category) as unique_categories FROM "
|
||||
"prompts WHERE category IS NOT NULL"
|
||||
)
|
||||
unique_categories = cursor.fetchone()['unique_categories']
|
||||
|
||||
unique_categories = cursor.fetchone()["unique_categories"]
|
||||
|
||||
cursor = conn.execute(
|
||||
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS NOT NULL"
|
||||
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS "
|
||||
"NOT NULL"
|
||||
)
|
||||
avg_rating = cursor.fetchone()['avg_rating']
|
||||
|
||||
avg_rating = cursor.fetchone()["avg_rating"]
|
||||
|
||||
# Get database file size
|
||||
db_size = os.path.getsize(self.db_path) if os.path.exists(self.db_path) else 0
|
||||
|
||||
db_size = (
|
||||
os.path.getsize(self.db_path) if os.path.exists(self.db_path) else 0
|
||||
)
|
||||
|
||||
return {
|
||||
'total_prompts': total_prompts,
|
||||
'unique_categories': unique_categories,
|
||||
'average_rating': round(avg_rating, 2) if avg_rating else None,
|
||||
'database_size_bytes': db_size,
|
||||
'database_path': os.path.abspath(self.db_path)
|
||||
"total_prompts": total_prompts,
|
||||
"unique_categories": unique_categories,
|
||||
"average_rating": round(avg_rating, 2) if avg_rating else None,
|
||||
"database_size_bytes": db_size,
|
||||
"database_path": os.path.abspath(self.db_path),
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error getting database info: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def backup_database(self, backup_path: str) -> bool:
|
||||
"""
|
||||
Create a backup of the database.
|
||||
|
||||
Write a consistent copy of the database to backup_path.
|
||||
|
||||
Uses the SQLite online backup API through this thread's connection,
|
||||
so rows that are committed but still sit in the -wal file are
|
||||
included and other threads may keep reading and writing meanwhile.
|
||||
|
||||
Args:
|
||||
backup_path: Path where the backup should be saved
|
||||
|
||||
backup_path: Path where the backup should be saved (overwritten)
|
||||
|
||||
Returns:
|
||||
bool: True if backup was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
import shutil
|
||||
shutil.copy2(self.db_path, backup_path)
|
||||
if os.path.exists(backup_path) and os.path.samefile(
|
||||
backup_path, self.db_path
|
||||
):
|
||||
self.logger.error("Refusing to back up the database onto itself")
|
||||
return False
|
||||
if os.path.exists(backup_path):
|
||||
os.remove(backup_path)
|
||||
source = self.get_connection()
|
||||
destination = sqlite3.connect(backup_path)
|
||||
try:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
destination.close()
|
||||
return True
|
||||
except Exception as e:
|
||||
except (OSError, sqlite3.Error) as e:
|
||||
self.logger.error(f"Error creating database backup: {e}")
|
||||
return False
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def verify_database_file(path: str) -> Tuple[bool, str]:
|
||||
"""
|
||||
Check that a file is a healthy prompts database before restoring it.
|
||||
|
||||
Opens the file read-only, runs PRAGMA integrity_check and confirms the
|
||||
prompts table has the columns every version of this extension needs.
|
||||
|
||||
Args:
|
||||
path: File to inspect
|
||||
|
||||
Returns:
|
||||
(ok, reason): reason is "ok" when the file is acceptable, otherwise
|
||||
a human-readable explanation. Never raises for bad input.
|
||||
"""
|
||||
if not os.path.isfile(path):
|
||||
return False, "File not found"
|
||||
try:
|
||||
with open(path, "rb") as fh:
|
||||
header = fh.read(len(SQLITE_HEADER))
|
||||
except OSError as e:
|
||||
return False, f"Cannot read file: {e}"
|
||||
if header != SQLITE_HEADER:
|
||||
return False, "Not a SQLite database file"
|
||||
|
||||
uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro"
|
||||
try:
|
||||
conn = sqlite3.connect(uri, uri=True)
|
||||
except sqlite3.Error as e:
|
||||
return False, f"Invalid SQLite database: {e}"
|
||||
try:
|
||||
row = conn.execute("PRAGMA integrity_check").fetchone()
|
||||
verdict = row[0] if row else "no result"
|
||||
if verdict != "ok":
|
||||
return False, f"Integrity check failed: {verdict}"
|
||||
table = conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'"
|
||||
).fetchone()
|
||||
if not table:
|
||||
return False, "Database does not contain a 'prompts' table"
|
||||
columns = {r[1] for r in conn.execute("PRAGMA table_info(prompts)")}
|
||||
missing = [c for c in REQUIRED_PROMPT_COLUMNS if c not in columns]
|
||||
if missing:
|
||||
return (
|
||||
False,
|
||||
"prompts table is missing required columns: " + ", ".join(missing),
|
||||
)
|
||||
return True, "ok"
|
||||
except sqlite3.Error as e:
|
||||
return False, f"Invalid SQLite database: {e}"
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _replace_contents_from(self, src_path: str) -> None:
|
||||
"""Copy every page of *src_path* into the live database.
|
||||
|
||||
Goes through SQLite's online backup API into this thread's normal
|
||||
connection, so SQLite takes the locks itself and the other
|
||||
PromptModel instances that share the file (nodes, image monitor,
|
||||
queue hook, API) simply see the new pages on their next statement.
|
||||
Nothing is swapped underneath open handles and no -wal/-shm file has
|
||||
to be deleted, which also keeps Windows happy. When SQLite refuses
|
||||
(page sizes differ while the target is in WAL mode) fall back to a
|
||||
file copy with every connection of this instance closed first.
|
||||
"""
|
||||
source = sqlite3.connect(
|
||||
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
|
||||
)
|
||||
try:
|
||||
try:
|
||||
source.backup(self.get_connection())
|
||||
return
|
||||
except sqlite3.OperationalError as e:
|
||||
self.logger.warning(
|
||||
f"Online restore failed ({e}); falling back to a file copy"
|
||||
)
|
||||
finally:
|
||||
source.close()
|
||||
self.close_all()
|
||||
shutil.copy2(src_path, self.db_path)
|
||||
for suffix in ("-wal", "-shm"):
|
||||
stale = self.db_path + suffix
|
||||
if os.path.exists(stale):
|
||||
os.remove(stale)
|
||||
|
||||
def restore_from_file(self, src_path: str) -> str:
|
||||
"""
|
||||
Replace the live database with the file at src_path.
|
||||
|
||||
The current database is first backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), the
|
||||
pages of ``src_path`` are copied into the live database through
|
||||
SQLite's backup API (visible to every open connection) and the schema
|
||||
migrations run on the restored content.
|
||||
|
||||
Callers should run verify_database_file(src_path) first.
|
||||
|
||||
Args:
|
||||
src_path: SQLite file to restore
|
||||
|
||||
Returns:
|
||||
Path of the backup taken of the previous database, or "" when
|
||||
there was no database to back up.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the previous database could not be backed up
|
||||
OSError: If the restored file could not be copied into place
|
||||
"""
|
||||
backup_path = ""
|
||||
if os.path.exists(self.db_path):
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_path = f"{self.db_path}.backup_{timestamp}"
|
||||
if not self.backup_database(backup_path):
|
||||
raise RuntimeError("Could not back up the current database")
|
||||
self.logger.info(f"Current database backed up to: {backup_path}")
|
||||
|
||||
self._replace_contents_from(src_path)
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.discard(self._schema_key())
|
||||
self._ensure_database_exists()
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
return backup_path
|
||||
|
||||
|
After Width: | Height: | Size: 75 KiB |
|
After Width: | Height: | Size: 74 KiB |
|
After Width: | Height: | Size: 76 KiB |
|
After Width: | Height: | Size: 75 KiB |
|
After Width: | Height: | Size: 78 KiB |
|
After Width: | Height: | Size: 102 KiB |
|
After Width: | Height: | Size: 104 KiB |
|
After Width: | Height: | Size: 79 KiB |
|
After Width: | Height: | Size: 92 KiB |
|
After Width: | Height: | Size: 94 KiB |
|
After Width: | Height: | Size: 79 KiB |
|
After Width: | Height: | Size: 100 KiB |
|
After Width: | Height: | Size: 101 KiB |
|
After Width: | Height: | Size: 80 KiB |
|
After Width: | Height: | Size: 90 KiB |
|
After Width: | Height: | Size: 91 KiB |
@@ -0,0 +1,71 @@
|
||||
# Shared ComfyAssets UI adoption
|
||||
|
||||
Status: approved direction; toolkit/testing foundation in progress.
|
||||
Product UI brainstorming comes before node changes.
|
||||
|
||||
Tracking: [PromptManager #186](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/186)
|
||||
and [KikoTools #81](https://github.com/ComfyAssets/ComfyUI-KikoTools/issues/81).
|
||||
|
||||
KikoTools is the validated model. The
|
||||
[ComfyAssets Toolkit](https://github.com/ComfyAssets/ComfyAssets-Toolkit) is the
|
||||
shared source; PromptManager proves the package boundaries before a later
|
||||
KikoTools adoption PR. Other custom nodes can become additional consumers.
|
||||
|
||||
## Approved sequence
|
||||
|
||||
1. Extract reusable controls and test infrastructure from the working KikoTools
|
||||
reference into the toolkit. Keep routes, node schemas and controllers local.
|
||||
2. Bring PromptManager's automated tests up to parity using the pinned toolkit.
|
||||
3. Brainstorm PromptManager node layouts, controls and behavior with the maintainer
|
||||
before changing its production UI. Select one bounded first consumer together.
|
||||
4. Validate shared distribution, themes, cleanup and compatibility through that
|
||||
consumer, then migrate incrementally. Test both packs together as appropriate.
|
||||
5. Backport toolkit adoption to KikoTools after PromptManager proves it, preserving
|
||||
the existing KikoTools appearance, behavior and workflow compatibility.
|
||||
|
||||
KikoTools is not being migrated again during this phase. The original
|
||||
[foundation mockup](https://github.com/ComfyAssets/ComfyUI-KikoTools/pull/82)
|
||||
records the approved visual direction and inspiration.
|
||||
|
||||
## PromptManager requirements
|
||||
|
||||
- Preserve existing node names, inputs/outputs, API routes, saved workflows,
|
||||
gallery behavior, and existing test/coverage gates.
|
||||
- Use shared control factories, editor slots, settings/help/window behavior and
|
||||
the ComfyUI adapter where applicable. Keep PromptManager controllers and HTTP
|
||||
behavior in this repository.
|
||||
- Both projects must look identical: one shared set of controls, typography,
|
||||
spacing, colors and interaction states. Follow ComfyUI light/dark mode with
|
||||
our own styling, a blue primary accent and strong red errors/destructive actions.
|
||||
Use compact buttons, pill switches, segmented selections and responsive rows.
|
||||
Review the pilot mockup before production styling. Existing `--pm-*` tokens
|
||||
are migration context, not an approved design to retain.
|
||||
- Theme standalone pages explicitly: `window.open` pages do not inherit the
|
||||
host's CSS variables. Test the agreed palette policy, live changes, fallbacks
|
||||
and persistence, including errors when host settings are unavailable.
|
||||
- Ship a pinned, reproducible browser distribution. Installation must work
|
||||
without npm or KikoTools, including when both packs ship different core versions.
|
||||
- Attribute [Pixaroma](https://github.com/pixaroma/ComfyUI-Pixaroma) as architectural
|
||||
inspiration. Do not copy/vendor his code, styles, or assets.
|
||||
|
||||
## First migration
|
||||
|
||||
With the maintainer, inventory the PromptManager nodes and gallery dialogs and
|
||||
choose one bounded first UI consumer after brainstorming. Record its current behavior with tests before changing it.
|
||||
Verify keyboard opening/closing, focus restoration, accessible labels, confirm
|
||||
and cancel behavior, network failures, and light/dark appearances. Keep the
|
||||
existing gallery and destructive-action browser workflows passing.
|
||||
|
||||
Build shared-source extraction and distribution checks before importing the core.
|
||||
Test both packs alone and together, scoped styling, isolated settings/storage,
|
||||
and independent cleanup. ComfyUI-facing changes require the real-host Classic,
|
||||
Nodes 2.0 and live-switching matrix from the foundation proposal. Standalone
|
||||
dialogs additionally run through PromptManager's existing browser harness.
|
||||
|
||||
## Delivery
|
||||
|
||||
Use small issue-linked PRs and TDD: verify each new behavioral test fails before
|
||||
implementing it, then run relevant suites and required CI checks. Include visual
|
||||
evidence for appearance changes. Migrate additional nodes/dialogs/editors only after
|
||||
the first consumer validates the package boundaries; do not rewrite all pages
|
||||
or change backend schemas as part of this initiative.
|
||||
|
After Width: | Height: | Size: 50 KiB |
|
After Width: | Height: | Size: 38 KiB |
|
After Width: | Height: | Size: 38 KiB |
|
After Width: | Height: | Size: 34 KiB |
|
After Width: | Height: | Size: 45 KiB |
@@ -3,22 +3,8 @@ PromptManager: Main custom node implementation that extends CLIPTextEncode
|
||||
with persistent prompt storage and search capabilities.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import webbrowser
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from utils.logging_config import get_logger
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
try:
|
||||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||||
@@ -35,23 +21,16 @@ except ImportError:
|
||||
InputTypeDict = dict
|
||||
|
||||
try:
|
||||
from .database.operations import PromptDatabase
|
||||
from .utils.comfyui_integration import get_comfyui_integration
|
||||
from .utils.image_monitor import ImageMonitor
|
||||
from .utils.prompt_tracker import PromptExecutionContext, PromptTracker
|
||||
from .prompt_manager_base import PromptManagerBase
|
||||
except ImportError:
|
||||
# For direct imports when not in a package
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from database.operations import PromptDatabase
|
||||
from utils.comfyui_integration import get_comfyui_integration
|
||||
from utils.image_monitor import ImageMonitor
|
||||
from utils.prompt_tracker import PromptExecutionContext, PromptTracker
|
||||
from prompt_manager_base import PromptManagerBase
|
||||
|
||||
|
||||
class PromptManager(ComfyNodeABC):
|
||||
class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
"""
|
||||
A ComfyUI custom node that functions like CLIPTextEncode but adds:
|
||||
- Persistent storage of all prompts in SQLite database
|
||||
@@ -61,17 +40,7 @@ class PromptManager(ComfyNodeABC):
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.logger = get_logger("prompt_manager.node")
|
||||
self.logger.debug("Initializing PromptManager node")
|
||||
|
||||
self.db = PromptDatabase()
|
||||
self.prompt_tracker = PromptTracker(self.db)
|
||||
self.image_monitor = ImageMonitor(self.db, self.prompt_tracker)
|
||||
self.comfyui_integration = get_comfyui_integration()
|
||||
|
||||
# Start image monitoring automatically
|
||||
self._start_gallery_system()
|
||||
self.logger.debug("PromptManager node initialization completed")
|
||||
super().__init__(logger_name="prompt_manager.node")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
@@ -82,7 +51,9 @@ class PromptManager(ComfyNodeABC):
|
||||
{
|
||||
"multiline": True,
|
||||
"dynamicPrompts": True,
|
||||
"tooltip": "The text prompt to be encoded and saved to database.",
|
||||
"tooltip": (
|
||||
"The text prompt to be encoded and saved to database."
|
||||
),
|
||||
},
|
||||
),
|
||||
"clip": (
|
||||
@@ -95,52 +66,77 @@ class PromptManager(ComfyNodeABC):
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Optional category for organizing prompts (e.g., 'landscapes', 'portraits')",
|
||||
"tooltip": (
|
||||
"Optional category for organizing prompts "
|
||||
"(e.g., 'landscapes', 'portraits')"
|
||||
),
|
||||
},
|
||||
),
|
||||
"tags": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Comma-separated tags for the prompt (e.g., 'anime, detailed, sunset')",
|
||||
"tooltip": (
|
||||
"Comma-separated tags for the prompt "
|
||||
"(e.g., 'anime, detailed, sunset')"
|
||||
),
|
||||
},
|
||||
),
|
||||
"search_text": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Search for past prompts containing this text",
|
||||
"tooltip": (
|
||||
"Search past prompts by text; add tag:name "
|
||||
'(or tag:"two words") to filter by tag'
|
||||
),
|
||||
},
|
||||
),
|
||||
"prepend_text": (
|
||||
IO.STRING,
|
||||
{
|
||||
"tooltip": "Text to prepend to the main prompt (connected STRING nodes will be added before the main text)"
|
||||
"tooltip": (
|
||||
"Text to prepend to the main prompt (connected STRING "
|
||||
"nodes will be added before the main text)"
|
||||
)
|
||||
},
|
||||
),
|
||||
"append_text": (
|
||||
IO.STRING,
|
||||
{
|
||||
"tooltip": "Text to append to the main prompt (connected STRING nodes will be added after the main text)"
|
||||
"tooltip": (
|
||||
"Text to append to the main prompt (connected STRING "
|
||||
"nodes will be added after the main text)"
|
||||
)
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
# The queued graph and this node's id: used to tell the positive
|
||||
# prompt of a run from negative-prompt PromptManager nodes
|
||||
"prompt": "PROMPT",
|
||||
"unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.CONDITIONING, IO.STRING)
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"A conditioning containing the embedded text used to guide the diffusion model.",
|
||||
"The final combined text string (with prepend/append applied) that was encoded.",
|
||||
"A conditioning containing the embedded text used to guide the diffusion "
|
||||
"model.",
|
||||
"The final combined text string (with prepend/append applied) that was "
|
||||
"encoded.",
|
||||
)
|
||||
FUNCTION = "encode"
|
||||
FUNCTION = "encode_prompt"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
|
||||
DESCRIPTION = (
|
||||
"Encodes a text prompt using a CLIP model into an embedding that can be used to guide "
|
||||
"the diffusion model towards generating specific images. Additionally saves all prompts "
|
||||
"Encodes a text prompt using a CLIP model into an embedding that can be used "
|
||||
"to guide the diffusion model towards generating specific images. "
|
||||
"Additionally saves all prompts "
|
||||
"to a local SQLite database with optional metadata for search and retrieval."
|
||||
)
|
||||
|
||||
def encode(
|
||||
def encode_prompt(
|
||||
self,
|
||||
clip,
|
||||
text: str,
|
||||
@@ -149,6 +145,8 @@ class PromptManager(ComfyNodeABC):
|
||||
search_text: str = "",
|
||||
prepend_text: str = "",
|
||||
append_text: str = "",
|
||||
prompt: Optional[Dict[str, Any]] = None,
|
||||
unique_id: Optional[str] = None,
|
||||
) -> Tuple[Any]:
|
||||
"""
|
||||
Encode the text prompt and save it to the database.
|
||||
@@ -163,28 +161,32 @@ class PromptManager(ComfyNodeABC):
|
||||
append_text: Text to append to the main prompt
|
||||
|
||||
Returns:
|
||||
Tuple containing the conditioning for the diffusion model and the final text string
|
||||
Tuple containing the conditioning for the diffusion model and the final
|
||||
text string
|
||||
|
||||
Raises:
|
||||
RuntimeError: If clip input is invalid
|
||||
"""
|
||||
# Combine prepend, main text, and append text
|
||||
final_text = ""
|
||||
parts = []
|
||||
if prepend_text and prepend_text.strip():
|
||||
final_text += prepend_text.strip() + " "
|
||||
final_text += text
|
||||
parts.append(prepend_text.strip())
|
||||
if text:
|
||||
parts.append(text)
|
||||
if append_text and append_text.strip():
|
||||
final_text += " " + append_text.strip()
|
||||
parts.append(append_text.strip())
|
||||
final_text = " ".join(parts)
|
||||
|
||||
# Inject LoRA trigger words if integration is enabled
|
||||
final_text = self._inject_lora_trigger_words(final_text)
|
||||
|
||||
# Use the combined text for encoding
|
||||
encoding_text = final_text
|
||||
|
||||
# For database storage, save the original main text with metadata about prepend/append
|
||||
# For database storage, save the original main text with metadata
|
||||
# about prepend/append
|
||||
storage_text = text
|
||||
|
||||
# Search functionality is now handled by the JavaScript UI
|
||||
# The search parameters are still available for backend processing if needed
|
||||
|
||||
# Validate CLIP model
|
||||
if clip is None:
|
||||
error_msg = (
|
||||
@@ -197,6 +199,7 @@ class PromptManager(ComfyNodeABC):
|
||||
|
||||
# Save prompt to database and set execution context for gallery tracking
|
||||
prompt_id = None
|
||||
extended_tags = []
|
||||
if storage_text and storage_text.strip():
|
||||
self.logger.debug(f"Processing prompt text: {storage_text[:100]}...")
|
||||
|
||||
@@ -208,35 +211,28 @@ class PromptManager(ComfyNodeABC):
|
||||
extended_tags.append(f"append:{append_text.strip()[:50]}")
|
||||
|
||||
try:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=storage_text.strip(), # Always strip whitespace
|
||||
prompt_id = self._track_prompt_execution(
|
||||
text=storage_text.strip(),
|
||||
encoding_text=encoding_text.strip(),
|
||||
category=category.strip() if category else None,
|
||||
tags=extended_tags if extended_tags else None,
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
"final_text": encoding_text.strip(),
|
||||
},
|
||||
prompt_graph=prompt,
|
||||
unique_id=unique_id,
|
||||
)
|
||||
|
||||
# Set current prompt for image tracking
|
||||
if prompt_id:
|
||||
execution_id = self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=encoding_text.strip(), # Use final combined text for tracking
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prompt_id": prompt_id,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
"final_text": encoding_text.strip(), # Store final combined text
|
||||
},
|
||||
)
|
||||
self.logger.debug(
|
||||
f"Set execution context: {execution_id} for prompt ID: {prompt_id}"
|
||||
)
|
||||
self.logger.debug(f"Tracked prompt ID: {prompt_id}")
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the encoding
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
# Already logged above, no need for additional print
|
||||
|
||||
# Perform standard CLIP text encoding using the combined text
|
||||
self.logger.debug(
|
||||
@@ -246,7 +242,7 @@ class PromptManager(ComfyNodeABC):
|
||||
conditioning = clip.encode_from_tokens_scheduled(tokens)
|
||||
|
||||
# Register with ComfyUI integration for standard metadata compatibility
|
||||
node_id = f"promptmanager_{int(time.time() * 1000)}" # Unique node ID
|
||||
node_id = f"promptmanager_{int(time.time() * 1000)}"
|
||||
self.comfyui_integration.register_prompt(
|
||||
node_id,
|
||||
encoding_text.strip(),
|
||||
@@ -259,250 +255,44 @@ class PromptManager(ComfyNodeABC):
|
||||
},
|
||||
)
|
||||
|
||||
self.logger.debug("CLIP encoding completed successfully")
|
||||
self.logger.info(f"CLIP encoding completed, text: {repr(encoding_text)[:80]}")
|
||||
return (conditioning, encoding_text)
|
||||
|
||||
def _save_prompt_to_database(
|
||||
self, text: str, category: Optional[str] = None, tags: Optional[list] = None
|
||||
) -> Optional[int]:
|
||||
"""
|
||||
Save the prompt to the SQLite database.
|
||||
|
||||
Args:
|
||||
text: The prompt text
|
||||
category: Optional category
|
||||
tags: List of tags
|
||||
|
||||
Returns:
|
||||
The prompt ID if saved successfully, None otherwise
|
||||
"""
|
||||
try:
|
||||
# Generate hash for duplicate detection
|
||||
prompt_hash = self._generate_hash(text)
|
||||
self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...")
|
||||
|
||||
# Check if prompt already exists
|
||||
existing = self.db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
self.logger.info(
|
||||
f"Found existing prompt with ID {existing['id']}, updating metadata"
|
||||
)
|
||||
# Update metadata if this is a duplicate with new info
|
||||
if any([category, tags]):
|
||||
self.db.update_prompt_metadata(
|
||||
prompt_id=existing["id"], category=category, tags=tags
|
||||
)
|
||||
self.logger.debug("Updated metadata for existing prompt")
|
||||
return existing["id"]
|
||||
|
||||
# Save new prompt
|
||||
self.logger.debug(
|
||||
f"Saving new prompt with category: {category}, tags: {tags}"
|
||||
)
|
||||
prompt_id = self.db.save_prompt(
|
||||
text=text, category=category, tags=tags, prompt_hash=prompt_hash
|
||||
)
|
||||
|
||||
if prompt_id:
|
||||
self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
|
||||
else:
|
||||
self.logger.warning("Failed to save prompt - no ID returned")
|
||||
|
||||
return prompt_id
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error saving prompt to database: {e}")
|
||||
# Already logged above, no need for additional print
|
||||
return None
|
||||
|
||||
def _generate_hash(self, text: str) -> str:
|
||||
"""
|
||||
Generate SHA256 hash for the prompt text.
|
||||
|
||||
Args:
|
||||
text: The prompt text to hash
|
||||
|
||||
Returns:
|
||||
Hexadecimal string representation of the SHA256 hash
|
||||
"""
|
||||
# Normalize text for consistent hashing (strip whitespace, normalize case)
|
||||
normalized_text = text.strip().lower()
|
||||
return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest()
|
||||
|
||||
def _parse_tags(self, tags_string: str) -> Optional[list]:
|
||||
"""
|
||||
Parse comma-separated tags string into a list.
|
||||
|
||||
Args:
|
||||
tags_string: Comma-separated string of tags
|
||||
|
||||
Returns:
|
||||
List of parsed tags, or None if no valid tags found
|
||||
"""
|
||||
if not tags_string or not tags_string.strip():
|
||||
return None
|
||||
|
||||
tags = [tag.strip() for tag in tags_string.split(",") if tag.strip()]
|
||||
return tags if tags else None
|
||||
|
||||
def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search for past prompts by text content.
|
||||
|
||||
Args:
|
||||
search_text: Text to search for in prompt database
|
||||
|
||||
Returns:
|
||||
List of matching prompt dictionaries with metadata
|
||||
"""
|
||||
try:
|
||||
if not search_text or not search_text.strip():
|
||||
return []
|
||||
|
||||
results = self.db.search_prompts(
|
||||
text=search_text.strip(),
|
||||
category=None,
|
||||
tags=None,
|
||||
rating_min=None,
|
||||
limit=50,
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error searching prompts: {e}")
|
||||
return []
|
||||
|
||||
def _open_web_interface(self):
|
||||
"""
|
||||
Open the web interface in the default browser.
|
||||
|
||||
Attempts to locate and open the web interface HTML file.
|
||||
Logs warnings if the interface is not properly configured.
|
||||
"""
|
||||
try:
|
||||
# Look for a web interface directory
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
web_dir = os.path.join(current_dir, "web_interface")
|
||||
|
||||
if os.path.exists(web_dir):
|
||||
# If web interface exists, try to start it
|
||||
index_path = os.path.join(web_dir, "index.html")
|
||||
if os.path.exists(index_path):
|
||||
webbrowser.open(f"file://{index_path}")
|
||||
self.logger.info("Web interface opened in browser")
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"Web interface directory found but no index.html. Please check {web_dir} for setup instructions"
|
||||
)
|
||||
else:
|
||||
self.logger.info(
|
||||
"Web interface not yet implemented. This feature will open a web-based prompt management interface when the web_interface directory is created."
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error opening web interface: {e}")
|
||||
|
||||
def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]:
|
||||
"""
|
||||
API method for JavaScript UI to search prompts.
|
||||
|
||||
Args:
|
||||
search_text: Text to search for in prompts
|
||||
|
||||
Returns:
|
||||
List of matching prompt dictionaries
|
||||
"""
|
||||
return self._search_prompts(search_text=search_text)
|
||||
|
||||
def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
API method for JavaScript UI to get recent prompts.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of recent prompts to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent prompt dictionaries ordered by creation time
|
||||
"""
|
||||
try:
|
||||
return self.db.get_recent_prompts(limit=limit)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error getting recent prompts: {e}")
|
||||
return []
|
||||
|
||||
def _start_gallery_system(self):
|
||||
"""
|
||||
Initialize and start the gallery monitoring system.
|
||||
|
||||
Starts the image monitor which watches for new generated images
|
||||
and links them to their source prompts in the database.
|
||||
"""
|
||||
try:
|
||||
self.logger.debug("Starting gallery system...")
|
||||
|
||||
# Start image monitoring
|
||||
self.image_monitor.start_monitoring()
|
||||
|
||||
self.logger.debug("Gallery system started successfully")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to start gallery system: {e}")
|
||||
self.logger.warning("Gallery features will be disabled")
|
||||
|
||||
def get_gallery_status(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get status of the gallery system.
|
||||
|
||||
Returns:
|
||||
Dictionary containing status information for image monitor and prompt tracker
|
||||
"""
|
||||
return {
|
||||
"image_monitor": self.image_monitor.get_status(),
|
||||
"prompt_tracker": self.prompt_tracker.get_status(),
|
||||
}
|
||||
|
||||
def cleanup_gallery_system(self):
|
||||
"""
|
||||
Clean up gallery system resources.
|
||||
|
||||
Stops image monitoring and releases associated resources.
|
||||
Called automatically during object destruction.
|
||||
"""
|
||||
try:
|
||||
if hasattr(self, "image_monitor"):
|
||||
self.image_monitor.stop_monitoring()
|
||||
self.logger.debug("Gallery system cleaned up")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error cleaning up gallery system: {e}")
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Cleanup when object is destroyed.
|
||||
|
||||
Ensures proper resource cleanup by stopping the gallery system.
|
||||
"""
|
||||
self.cleanup_gallery_system()
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, text="", category="", tags="", search_text="",
|
||||
prepend_text="", append_text="", **kwargs):
|
||||
def IS_CHANGED(
|
||||
cls,
|
||||
clip,
|
||||
text="",
|
||||
category="",
|
||||
tags="",
|
||||
search_text="",
|
||||
prepend_text="",
|
||||
append_text="",
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
ComfyUI method to determine if node needs re-execution.
|
||||
|
||||
This method now properly tracks input changes to avoid unnecessary
|
||||
re-execution while still ensuring prompts are saved when inputs change.
|
||||
|
||||
Returns:
|
||||
A hash of the input values that changes when any input changes
|
||||
|
||||
Returns a hash of input values that affect the conditioning output or
|
||||
the saved prompt's metadata. This enables proper branch execution -
|
||||
only re-execute when inputs change.
|
||||
"""
|
||||
# Create a hash of all the text inputs that affect the output
|
||||
# This ensures the node only re-executes when inputs actually change
|
||||
import hashlib
|
||||
|
||||
# Combine all text inputs that affect the conditioning output
|
||||
combined = f"{text}|{category}|{tags}|{prepend_text}|{append_text}"
|
||||
|
||||
# Return a hash that will change when inputs change
|
||||
# Note: We don't include search_text as it doesn't affect the output conditioning
|
||||
|
||||
combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}"
|
||||
return hashlib.sha256(combined.encode()).hexdigest()
|
||||
|
||||
|
||||
class PromptManagerFocused(PromptManager):
|
||||
"""The editor-first UI uses the existing encoding and storage contract."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
inputs = super().INPUT_TYPES()
|
||||
for name in ("prepend_text", "append_text"):
|
||||
kind, options = inputs["optional"][name]
|
||||
inputs["optional"][name] = (
|
||||
kind,
|
||||
{**options, "default": "", "forceInput": True},
|
||||
)
|
||||
return inputs
|
||||
|
||||
@@ -0,0 +1,314 @@
|
||||
"""
|
||||
PromptManagerBase: Shared logic for PromptManager node variants.
|
||||
|
||||
Provides database initialization, prompt saving, hashing, tag parsing,
|
||||
search, gallery system management, and cleanup — extracted from the
|
||||
duplicate code in prompt_manager.py and prompt_manager_text.py.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
try:
|
||||
from .database.operations import PromptDatabase
|
||||
from .utils.comfyui_integration import get_comfyui_integration
|
||||
from .utils.image_monitor import get_image_monitor
|
||||
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.usage_tracking import PENDING_FIRST_USE
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from database.operations import PromptDatabase
|
||||
from utils.comfyui_integration import get_comfyui_integration
|
||||
from utils.image_monitor import get_image_monitor
|
||||
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.usage_tracking import PENDING_FIRST_USE
|
||||
|
||||
|
||||
class PromptManagerBase:
|
||||
"""Mixin providing shared prompt management logic for ComfyUI nodes.
|
||||
|
||||
Handles database connection, prompt saving with deduplication,
|
||||
tag parsing, search, gallery system lifecycle, and cleanup.
|
||||
Subclasses only need to define ComfyUI-specific class attributes
|
||||
(INPUT_TYPES, RETURN_TYPES, FUNCTION) and their execution method.
|
||||
"""
|
||||
|
||||
def __init__(self, logger_name: str = "prompt_manager.node"):
|
||||
self.logger = get_logger(logger_name)
|
||||
self.logger.debug(f"Initializing {self.__class__.__name__} node")
|
||||
|
||||
self.db = PromptDatabase()
|
||||
self.prompt_tracker = get_prompt_tracker(self.db)
|
||||
self.image_monitor = get_image_monitor(self.db, self.prompt_tracker)
|
||||
self.comfyui_integration = get_comfyui_integration()
|
||||
|
||||
self._start_gallery_system()
|
||||
self.logger.debug(f"{self.__class__.__name__} node initialization completed")
|
||||
|
||||
def _save_prompt_to_database(
|
||||
self,
|
||||
text: str,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[list] = None,
|
||||
count_run: bool = False,
|
||||
) -> Optional[int]:
|
||||
"""Save the prompt to the SQLite database.
|
||||
|
||||
Runs are normally counted by the queue hook (utils/usage_tracking.py), so a
|
||||
save only counts one when asked to (batch items whose text the hook could not
|
||||
see) or when the hook queued this prompt before it existed.
|
||||
|
||||
Args:
|
||||
text: The prompt text
|
||||
category: Optional category
|
||||
tags: List of tags
|
||||
count_run: Count this execution as a run of the prompt
|
||||
|
||||
Returns:
|
||||
The prompt ID if saved successfully, None otherwise
|
||||
"""
|
||||
try:
|
||||
prompt_hash = self._generate_hash(text)
|
||||
self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...")
|
||||
|
||||
existing = self.db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
self.logger.info(
|
||||
f"Found existing prompt with ID {existing['id']}, updating metadata"
|
||||
)
|
||||
if any([category, tags]):
|
||||
self.db.update_prompt_metadata(
|
||||
prompt_id=existing["id"], category=category, tags=tags
|
||||
)
|
||||
self.logger.debug("Updated metadata for existing prompt")
|
||||
if count_run:
|
||||
self._record_use(existing["id"])
|
||||
return existing["id"]
|
||||
|
||||
self.logger.debug(
|
||||
f"Saving new prompt with category: {category}, tags: {tags}"
|
||||
)
|
||||
prompt_id = self.db.save_prompt(
|
||||
text=text, category=category, tags=tags, prompt_hash=prompt_hash
|
||||
)
|
||||
|
||||
if prompt_id:
|
||||
self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
|
||||
pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
|
||||
runs = pending.consume(prompt_hash) + (1 if count_run else 0)
|
||||
if runs:
|
||||
self._record_use(prompt_id, times=runs)
|
||||
else:
|
||||
self.logger.warning("Failed to save prompt - no ID returned")
|
||||
|
||||
return prompt_id
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error saving prompt to database: {e}")
|
||||
return None
|
||||
|
||||
def _prompt_role(
|
||||
self, prompt_graph: Any, unique_id: Any, text: str
|
||||
) -> Dict[str, bool]:
|
||||
"""How this node's execution relates to the run, from the queued graph.
|
||||
|
||||
Returns flags:
|
||||
is_run_prompt: feeds a sampler's positive input (images belong to it)
|
||||
seen_at_queue: the queue hook resolved exactly this text, so it already
|
||||
counted the run and the image's metadata can identify the prompt
|
||||
known: the graph and node id were available
|
||||
"""
|
||||
if not isinstance(prompt_graph, dict) or unique_id is None:
|
||||
return {"is_run_prompt": True, "seen_at_queue": False, "known": False}
|
||||
node_id = str(unique_id)
|
||||
resolved = resolve_text(prompt_graph, node_id)
|
||||
return {
|
||||
"is_run_prompt": node_id in run_prompt_nodes(prompt_graph),
|
||||
"seen_at_queue": resolved is not None
|
||||
and generate_prompt_hash(resolved) == generate_prompt_hash(text),
|
||||
"known": True,
|
||||
}
|
||||
|
||||
def _track_prompt_execution(
|
||||
self,
|
||||
text: str,
|
||||
encoding_text: str,
|
||||
category: Optional[str],
|
||||
tags: Optional[list],
|
||||
additional_data: Dict[str, Any],
|
||||
prompt_graph: Any = None,
|
||||
unique_id: Any = None,
|
||||
) -> Optional[int]:
|
||||
"""Save the prompt and, for the run's positive prompt, register it for linking.
|
||||
|
||||
Negative-prompt nodes are saved but never become the prompt their images are
|
||||
linked to. Text the queue hook could resolve is counted there and linked from
|
||||
the image's own metadata; anything else (PromptSearchList batches, unknown
|
||||
text nodes, or a resolution that doesn't match) is counted here and pushed
|
||||
onto the FIFO queue. Without graph context the pre-3.2.4 behaviour is kept.
|
||||
"""
|
||||
role = self._prompt_role(prompt_graph, unique_id, text)
|
||||
batch_item = (
|
||||
role["known"] and role["is_run_prompt"] and not role["seen_at_queue"]
|
||||
)
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=text, category=category, tags=tags, count_run=batch_item
|
||||
)
|
||||
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"],
|
||||
)
|
||||
return prompt_id
|
||||
|
||||
def _record_use(self, prompt_id: int, times: int = 1) -> None:
|
||||
"""Count runs; a tracking failure must never block generation."""
|
||||
try:
|
||||
self.db.record_prompt_use(prompt_id, times=times)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not record use of prompt {prompt_id}: {e}")
|
||||
|
||||
def _inject_lora_trigger_words(self, text: str) -> str:
|
||||
"""Append LoRA trigger words if integration is enabled.
|
||||
|
||||
Returns the text unchanged if the integration is disabled or
|
||||
LoraManager is not installed.
|
||||
"""
|
||||
try:
|
||||
from .py.config import IntegrationConfig
|
||||
except ImportError:
|
||||
try:
|
||||
from py.config import IntegrationConfig
|
||||
except ImportError:
|
||||
return text
|
||||
|
||||
if (
|
||||
not IntegrationConfig.LORA_MANAGER_ENABLED
|
||||
or not IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED
|
||||
):
|
||||
return text
|
||||
|
||||
try:
|
||||
from .py.lora_utils import get_trigger_cache, inject_trigger_words
|
||||
except ImportError:
|
||||
try:
|
||||
from py.lora_utils import get_trigger_cache, inject_trigger_words
|
||||
except ImportError:
|
||||
return text
|
||||
|
||||
cache = get_trigger_cache()
|
||||
|
||||
# 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)
|
||||
if injected:
|
||||
self.logger.info(f"Injected trigger words: {', '.join(injected)}")
|
||||
return modified
|
||||
|
||||
def _generate_hash(self, text: str) -> str:
|
||||
"""The prompt's deduplication hash (see utils.hashing.generate_prompt_hash)."""
|
||||
return generate_prompt_hash(text)
|
||||
|
||||
def _parse_tags(self, tags_string: str) -> Optional[list]:
|
||||
"""Parse comma-separated tags string into a list.
|
||||
|
||||
Args:
|
||||
tags_string: Comma-separated string of tags
|
||||
|
||||
Returns:
|
||||
List of parsed tags, or None if no valid tags found
|
||||
"""
|
||||
if not tags_string or not tags_string.strip():
|
||||
return None
|
||||
|
||||
tags = [tag.strip() for tag in tags_string.split(",") if tag.strip()]
|
||||
return tags if tags else None
|
||||
|
||||
def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]:
|
||||
"""Search for past prompts by text content.
|
||||
|
||||
Args:
|
||||
search_text: Text to search for in prompt database
|
||||
|
||||
Returns:
|
||||
List of matching prompt dictionaries with metadata
|
||||
"""
|
||||
try:
|
||||
if not search_text or not search_text.strip():
|
||||
return []
|
||||
|
||||
try:
|
||||
from .py.config import PromptManagerConfig
|
||||
|
||||
max_results = PromptManagerConfig.MAX_SEARCH_RESULTS
|
||||
except Exception:
|
||||
max_results = 100
|
||||
|
||||
results = self.db.search_prompts(
|
||||
text=search_text.strip(),
|
||||
category=None,
|
||||
tags=None,
|
||||
rating_min=None,
|
||||
limit=max_results,
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error searching prompts: {e}")
|
||||
return []
|
||||
|
||||
def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]:
|
||||
"""API method for JavaScript UI to search prompts."""
|
||||
return self._search_prompts(search_text=search_text)
|
||||
|
||||
def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]:
|
||||
"""API method for JavaScript UI to get recent prompts."""
|
||||
try:
|
||||
return self.db.get_recent_prompts(limit=limit)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error getting recent prompts: {e}")
|
||||
return []
|
||||
|
||||
def _start_gallery_system(self):
|
||||
"""Initialize and start the gallery monitoring system."""
|
||||
try:
|
||||
self.logger.debug("Starting gallery system...")
|
||||
self.image_monitor.start_monitoring()
|
||||
self.logger.debug("Gallery system started successfully")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to start gallery system: {e}")
|
||||
self.logger.warning("Gallery features will be disabled")
|
||||
|
||||
def get_gallery_status(self) -> Dict[str, Any]:
|
||||
"""Get status of the gallery system."""
|
||||
return {
|
||||
"image_monitor": self.image_monitor.get_status(),
|
||||
"prompt_tracker": self.prompt_tracker.get_status(),
|
||||
}
|
||||
|
||||
def cleanup_gallery_system(self):
|
||||
"""Clean up gallery system resources.
|
||||
|
||||
Note: The image monitor is a singleton shared across all node instances,
|
||||
so we intentionally do NOT stop it here. It should run for the entire
|
||||
ComfyUI session. Stopping it would kill monitoring for all nodes.
|
||||
"""
|
||||
self.logger.debug("Gallery system cleaned up (monitor continues running)")
|
||||
@@ -3,22 +3,8 @@ PromptManagerText: A text-only version of PromptManager that outputs STRING
|
||||
without CLIP encoding, while maintaining all database and search features.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import webbrowser
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from utils.logging_config import get_logger
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
try:
|
||||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||||
@@ -35,23 +21,16 @@ except ImportError:
|
||||
InputTypeDict = dict
|
||||
|
||||
try:
|
||||
from .database.operations import PromptDatabase
|
||||
from .utils.comfyui_integration import get_comfyui_integration
|
||||
from .utils.image_monitor import ImageMonitor
|
||||
from .utils.prompt_tracker import PromptExecutionContext, PromptTracker
|
||||
from .prompt_manager_base import PromptManagerBase
|
||||
except ImportError:
|
||||
# For direct imports when not in a package
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from database.operations import PromptDatabase
|
||||
from utils.comfyui_integration import get_comfyui_integration
|
||||
from utils.image_monitor import ImageMonitor
|
||||
from utils.prompt_tracker import PromptExecutionContext, PromptTracker
|
||||
from prompt_manager_base import PromptManagerBase
|
||||
|
||||
|
||||
class PromptManagerText(ComfyNodeABC):
|
||||
class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
"""
|
||||
A ComfyUI custom node that provides all PromptManager features but outputs
|
||||
only a STRING without CLIP encoding. Includes:
|
||||
@@ -63,17 +42,7 @@ class PromptManagerText(ComfyNodeABC):
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.logger = get_logger("prompt_manager_text.node")
|
||||
self.logger.debug("Initializing PromptManagerText node")
|
||||
|
||||
self.db = PromptDatabase()
|
||||
self.prompt_tracker = PromptTracker(self.db)
|
||||
self.image_monitor = ImageMonitor(self.db, self.prompt_tracker)
|
||||
self.comfyui_integration = get_comfyui_integration()
|
||||
|
||||
# Start image monitoring automatically
|
||||
self._start_gallery_system()
|
||||
self.logger.debug("PromptManagerText node initialization completed")
|
||||
super().__init__(logger_name="prompt_manager_text.node")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
@@ -84,7 +53,9 @@ class PromptManagerText(ComfyNodeABC):
|
||||
{
|
||||
"multiline": True,
|
||||
"dynamicPrompts": True,
|
||||
"tooltip": "The text prompt to be processed and saved to database.",
|
||||
"tooltip": (
|
||||
"The text prompt to be processed and saved to database."
|
||||
),
|
||||
},
|
||||
)
|
||||
},
|
||||
@@ -93,50 +64,74 @@ class PromptManagerText(ComfyNodeABC):
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Optional category for organizing prompts (e.g., 'landscapes', 'portraits')",
|
||||
"tooltip": (
|
||||
"Optional category for organizing prompts "
|
||||
"(e.g., 'landscapes', 'portraits')"
|
||||
),
|
||||
},
|
||||
),
|
||||
"tags": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Comma-separated tags for the prompt (e.g., 'anime, detailed, sunset')",
|
||||
"tooltip": (
|
||||
"Comma-separated tags for the prompt "
|
||||
"(e.g., 'anime, detailed, sunset')"
|
||||
),
|
||||
},
|
||||
),
|
||||
"search_text": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Search for past prompts containing this text",
|
||||
"tooltip": (
|
||||
"Search past prompts by text; add tag:name "
|
||||
'(or tag:"two words") to filter by tag'
|
||||
),
|
||||
},
|
||||
),
|
||||
"prepend_text": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Text to prepend to the main prompt (connected STRING nodes will be added before the main text)",
|
||||
"tooltip": (
|
||||
"Text to prepend to the main prompt (connected STRING "
|
||||
"nodes will be added before the main text)"
|
||||
),
|
||||
},
|
||||
),
|
||||
"append_text": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Text to append to the main prompt (connected STRING nodes will be added after the main text)",
|
||||
"tooltip": (
|
||||
"Text to append to the main prompt (connected STRING "
|
||||
"nodes will be added after the main text)"
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
# The queued graph and this node's id: used to tell the positive
|
||||
# prompt of a run from negative-prompt PromptManager nodes
|
||||
"prompt": "PROMPT",
|
||||
"unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.STRING,)
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"The final combined text string (with prepend/append applied) ready for use in other nodes.",
|
||||
"The final combined text string (with prepend/append applied) ready for use "
|
||||
"in other nodes.",
|
||||
)
|
||||
FUNCTION = "process_text"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
|
||||
DESCRIPTION = (
|
||||
"Processes and manages text prompts with database storage and search capabilities. "
|
||||
"Outputs a plain STRING that can be used with any node that accepts text input. "
|
||||
"Includes all PromptManager features: categorization, tagging, search, and prepend/append functionality."
|
||||
"Processes and manages text prompts with database storage and search "
|
||||
"capabilities. Outputs a plain STRING that can be used with any node that "
|
||||
"accepts text input. Includes all PromptManager features: categorization, "
|
||||
"tagging, search, and prepend/append functionality."
|
||||
)
|
||||
|
||||
def process_text(
|
||||
@@ -147,6 +142,8 @@ class PromptManagerText(ComfyNodeABC):
|
||||
search_text: str = "",
|
||||
prepend_text: str = "",
|
||||
append_text: str = "",
|
||||
prompt: Optional[Dict[str, Any]] = None,
|
||||
unique_id: Optional[str] = None,
|
||||
) -> Tuple[str]:
|
||||
"""
|
||||
Process the text prompt and save it to the database.
|
||||
@@ -163,21 +160,25 @@ class PromptManagerText(ComfyNodeABC):
|
||||
Tuple containing the final processed text string
|
||||
"""
|
||||
# Combine prepend, main text, and append text
|
||||
final_text = ""
|
||||
parts = []
|
||||
if prepend_text and prepend_text.strip():
|
||||
final_text += prepend_text.strip() + " "
|
||||
final_text += text
|
||||
parts.append(prepend_text.strip())
|
||||
if text:
|
||||
parts.append(text)
|
||||
if append_text and append_text.strip():
|
||||
final_text += " " + append_text.strip()
|
||||
parts.append(append_text.strip())
|
||||
final_text = " ".join(parts)
|
||||
|
||||
# For database storage, save the original main text with metadata about prepend/append
|
||||
# Inject LoRA trigger words if integration is enabled
|
||||
final_text = self._inject_lora_trigger_words(final_text)
|
||||
|
||||
# For database storage, save the original main text with metadata
|
||||
# about prepend/append
|
||||
storage_text = text
|
||||
|
||||
# Search functionality is now handled by the JavaScript UI
|
||||
# The search parameters are still available for backend processing if needed
|
||||
|
||||
# Save prompt to database and set execution context for gallery tracking
|
||||
prompt_id = None
|
||||
extended_tags = []
|
||||
if storage_text and storage_text.strip():
|
||||
self.logger.debug(f"Processing prompt text: {storage_text[:100]}...")
|
||||
|
||||
@@ -189,36 +190,30 @@ class PromptManagerText(ComfyNodeABC):
|
||||
extended_tags.append(f"append:{append_text.strip()[:50]}")
|
||||
|
||||
try:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=storage_text.strip(), # Always strip whitespace
|
||||
prompt_id = self._track_prompt_execution(
|
||||
text=storage_text.strip(),
|
||||
encoding_text=final_text.strip(),
|
||||
category=category.strip() if category else None,
|
||||
tags=extended_tags if extended_tags else None,
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
},
|
||||
prompt_graph=prompt,
|
||||
unique_id=unique_id,
|
||||
)
|
||||
|
||||
# Set current prompt for image tracking
|
||||
if prompt_id:
|
||||
execution_id = self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=final_text.strip(), # Use final combined text for tracking
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prompt_id": prompt_id,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
},
|
||||
)
|
||||
self.logger.debug(
|
||||
f"Set execution context: {execution_id} for prompt ID: {prompt_id}"
|
||||
)
|
||||
self.logger.debug(f"Tracked prompt ID: {prompt_id}")
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the processing
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
|
||||
# Register with ComfyUI integration for standard metadata compatibility
|
||||
node_id = f"promptmanagertext_{int(time.time() * 1000)}" # Unique node ID
|
||||
node_id = f"promptmanagertext_{int(time.time() * 1000)}"
|
||||
self.comfyui_integration.register_prompt(
|
||||
node_id,
|
||||
final_text.strip(),
|
||||
@@ -234,246 +229,25 @@ class PromptManagerText(ComfyNodeABC):
|
||||
self.logger.debug(f"Text processing completed: {final_text[:100]}...")
|
||||
return (final_text,)
|
||||
|
||||
def _save_prompt_to_database(
|
||||
self, text: str, category: Optional[str] = None, tags: Optional[list] = None
|
||||
) -> Optional[int]:
|
||||
"""
|
||||
Save the prompt to the SQLite database.
|
||||
|
||||
Args:
|
||||
text: The prompt text
|
||||
category: Optional category
|
||||
tags: List of tags
|
||||
|
||||
Returns:
|
||||
The prompt ID if saved successfully, None otherwise
|
||||
"""
|
||||
try:
|
||||
# Generate hash for duplicate detection
|
||||
prompt_hash = self._generate_hash(text)
|
||||
self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...")
|
||||
|
||||
# Check if prompt already exists
|
||||
existing = self.db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
self.logger.info(
|
||||
f"Found existing prompt with ID {existing['id']}, updating metadata"
|
||||
)
|
||||
# Update metadata if this is a duplicate with new info
|
||||
if any([category, tags]):
|
||||
self.db.update_prompt_metadata(
|
||||
prompt_id=existing["id"], category=category, tags=tags
|
||||
)
|
||||
self.logger.debug("Updated metadata for existing prompt")
|
||||
return existing["id"]
|
||||
|
||||
# Save new prompt
|
||||
self.logger.debug(
|
||||
f"Saving new prompt with category: {category}, tags: {tags}"
|
||||
)
|
||||
prompt_id = self.db.save_prompt(
|
||||
text=text, category=category, tags=tags, prompt_hash=prompt_hash
|
||||
)
|
||||
|
||||
if prompt_id:
|
||||
self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
|
||||
else:
|
||||
self.logger.warning("Failed to save prompt - no ID returned")
|
||||
|
||||
return prompt_id
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error saving prompt to database: {e}")
|
||||
return None
|
||||
|
||||
def _generate_hash(self, text: str) -> str:
|
||||
"""
|
||||
Generate SHA256 hash for the prompt text.
|
||||
|
||||
Args:
|
||||
text: The prompt text to hash
|
||||
|
||||
Returns:
|
||||
Hexadecimal string representation of the SHA256 hash
|
||||
"""
|
||||
# Normalize text for consistent hashing (strip whitespace, normalize case)
|
||||
normalized_text = text.strip().lower()
|
||||
return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest()
|
||||
|
||||
def _parse_tags(self, tags_string: str) -> Optional[list]:
|
||||
"""
|
||||
Parse comma-separated tags string into a list.
|
||||
|
||||
Args:
|
||||
tags_string: Comma-separated string of tags
|
||||
|
||||
Returns:
|
||||
List of parsed tags, or None if no valid tags found
|
||||
"""
|
||||
if not tags_string or not tags_string.strip():
|
||||
return None
|
||||
|
||||
tags = [tag.strip() for tag in tags_string.split(",") if tag.strip()]
|
||||
return tags if tags else None
|
||||
|
||||
def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search for past prompts by text content.
|
||||
|
||||
Args:
|
||||
search_text: Text to search for in prompt database
|
||||
|
||||
Returns:
|
||||
List of matching prompt dictionaries with metadata
|
||||
"""
|
||||
try:
|
||||
if not search_text or not search_text.strip():
|
||||
return []
|
||||
|
||||
results = self.db.search_prompts(
|
||||
text=search_text.strip(),
|
||||
category=None,
|
||||
tags=None,
|
||||
rating_min=None,
|
||||
limit=50,
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error searching prompts: {e}")
|
||||
return []
|
||||
|
||||
def _open_web_interface(self):
|
||||
"""
|
||||
Open the web interface in the default browser.
|
||||
|
||||
Attempts to locate and open the web interface HTML file.
|
||||
Logs warnings if the interface is not properly configured.
|
||||
"""
|
||||
try:
|
||||
# Look for a web interface directory
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
web_dir = os.path.join(current_dir, "web_interface")
|
||||
|
||||
if os.path.exists(web_dir):
|
||||
# If web interface exists, try to start it
|
||||
index_path = os.path.join(web_dir, "index.html")
|
||||
if os.path.exists(index_path):
|
||||
webbrowser.open(f"file://{index_path}")
|
||||
self.logger.info("Web interface opened in browser")
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"Web interface directory found but no index.html. Please check {web_dir} for setup instructions"
|
||||
)
|
||||
else:
|
||||
self.logger.info(
|
||||
"Web interface not yet implemented. This feature will open a web-based prompt management interface when the web_interface directory is created."
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error opening web interface: {e}")
|
||||
|
||||
def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]:
|
||||
"""
|
||||
API method for JavaScript UI to search prompts.
|
||||
|
||||
Args:
|
||||
search_text: Text to search for in prompts
|
||||
|
||||
Returns:
|
||||
List of matching prompt dictionaries
|
||||
"""
|
||||
return self._search_prompts(search_text=search_text)
|
||||
|
||||
def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
API method for JavaScript UI to get recent prompts.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of recent prompts to retrieve
|
||||
|
||||
Returns:
|
||||
List of recent prompt dictionaries ordered by creation time
|
||||
"""
|
||||
try:
|
||||
return self.db.get_recent_prompts(limit=limit)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error getting recent prompts: {e}")
|
||||
return []
|
||||
|
||||
def _start_gallery_system(self):
|
||||
"""
|
||||
Initialize and start the gallery monitoring system.
|
||||
|
||||
Starts the image monitor which watches for new generated images
|
||||
and links them to their source prompts in the database.
|
||||
"""
|
||||
try:
|
||||
self.logger.debug("Starting gallery system...")
|
||||
|
||||
# Start image monitoring
|
||||
self.image_monitor.start_monitoring()
|
||||
|
||||
self.logger.debug("Gallery system started successfully")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to start gallery system: {e}")
|
||||
self.logger.warning("Gallery features will be disabled")
|
||||
|
||||
def get_gallery_status(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get status of the gallery system.
|
||||
|
||||
Returns:
|
||||
Dictionary containing status information for image monitor and prompt tracker
|
||||
"""
|
||||
return {
|
||||
"image_monitor": self.image_monitor.get_status(),
|
||||
"prompt_tracker": self.prompt_tracker.get_status(),
|
||||
}
|
||||
|
||||
def cleanup_gallery_system(self):
|
||||
"""
|
||||
Clean up gallery system resources.
|
||||
|
||||
Stops image monitoring and releases associated resources.
|
||||
Called automatically during object destruction.
|
||||
"""
|
||||
try:
|
||||
if hasattr(self, "image_monitor"):
|
||||
self.image_monitor.stop_monitoring()
|
||||
self.logger.debug("Gallery system cleaned up")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error cleaning up gallery system: {e}")
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Cleanup when object is destroyed.
|
||||
|
||||
Ensures proper resource cleanup by stopping the gallery system.
|
||||
"""
|
||||
self.cleanup_gallery_system()
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, text="", category="", tags="", search_text="",
|
||||
prepend_text="", append_text="", **kwargs):
|
||||
def IS_CHANGED(
|
||||
cls,
|
||||
text="",
|
||||
category="",
|
||||
tags="",
|
||||
search_text="",
|
||||
prepend_text="",
|
||||
append_text="",
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
ComfyUI method to determine if node needs re-execution.
|
||||
|
||||
This method now properly tracks input changes to avoid unnecessary
|
||||
re-execution while still ensuring prompts are saved when inputs change.
|
||||
|
||||
Returns:
|
||||
A hash of the input values that changes when any input changes
|
||||
|
||||
Returns a hash of input values that affect the text output or the
|
||||
saved prompt's metadata. This enables proper branch execution - only
|
||||
re-execute when inputs change.
|
||||
"""
|
||||
# Create a hash of all the text inputs that affect the output
|
||||
# This ensures the node only re-executes when inputs actually change
|
||||
import hashlib
|
||||
|
||||
# Combine all text inputs that affect the output
|
||||
combined = f"{text}|{category}|{tags}|{prepend_text}|{append_text}"
|
||||
|
||||
# Return a hash that will change when inputs change
|
||||
# Note: We don't include search_text as it doesn't affect the output
|
||||
|
||||
combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}"
|
||||
return hashlib.sha256(combined.encode()).hexdigest()
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
"""
|
||||
PromptSearchList: A ComfyUI node that searches prompts and outputs results as a list.
|
||||
|
||||
This node enables batch processing workflows by outputting prompt texts with
|
||||
OUTPUT_IS_LIST=True, allowing direct connection to nodes that accept list inputs.
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
try:
|
||||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||||
except ImportError:
|
||||
|
||||
class ComfyNodeABC:
|
||||
pass
|
||||
|
||||
class IO:
|
||||
STRING = "STRING"
|
||||
INT = "INT"
|
||||
|
||||
InputTypeDict = dict
|
||||
|
||||
try:
|
||||
from .database.operations import PromptDatabase
|
||||
except ImportError:
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from database.operations import PromptDatabase
|
||||
|
||||
|
||||
class PromptSearchList(ComfyNodeABC):
|
||||
"""
|
||||
A ComfyUI node that searches the prompt database and outputs matching
|
||||
prompts as a list for batch processing workflows.
|
||||
|
||||
Features:
|
||||
- Search by text content
|
||||
- Filter by category, tags, and minimum rating
|
||||
- Outputs as OUTPUT_IS_LIST for batch node compatibility
|
||||
- Read-only operation (no database writes)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.logger = get_logger("prompt_search_list.node")
|
||||
self.logger.debug("Initializing PromptSearchList node")
|
||||
self.db = PromptDatabase()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
return {
|
||||
"required": {},
|
||||
"optional": {
|
||||
"text": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Search text to match against prompt content",
|
||||
},
|
||||
),
|
||||
"category": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Filter by specific category",
|
||||
},
|
||||
),
|
||||
"tags": (
|
||||
IO.STRING,
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": (
|
||||
"Comma-separated list of tags to filter by (partial match)"
|
||||
),
|
||||
},
|
||||
),
|
||||
"min_rating": (
|
||||
IO.INT,
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 5,
|
||||
"tooltip": "Minimum rating (0-5) to include in results",
|
||||
},
|
||||
),
|
||||
"limit": (
|
||||
IO.INT,
|
||||
{
|
||||
"default": 50,
|
||||
"min": 1,
|
||||
"max": 1000,
|
||||
"tooltip": "Maximum number of results to return",
|
||||
},
|
||||
),
|
||||
"skip_multipart": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": (
|
||||
"Skip prompts containing Clip_1/Clip_2/etc. "
|
||||
"multi-part markers"
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.STRING, IO.STRING, IO.INT)
|
||||
RETURN_NAMES = ("prompts", "preview", "count")
|
||||
OUTPUT_IS_LIST = (True, False, False)
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"List of prompt texts matching the search criteria.",
|
||||
"Preview of all found prompts (for display).",
|
||||
"Number of prompts found.",
|
||||
)
|
||||
FUNCTION = "search"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
|
||||
DESCRIPTION = (
|
||||
"Searches the prompt database and outputs matching prompts as a list. "
|
||||
"Use this node to retrieve stored prompts for batch processing workflows. "
|
||||
"'prompts' output sends each prompt individually to downstream nodes (batch). "
|
||||
"'preview' output shows all found prompts as one text block. "
|
||||
"'count' output shows how many prompts were found."
|
||||
)
|
||||
|
||||
def search(
|
||||
self,
|
||||
text: str = "",
|
||||
category: str = "",
|
||||
tags: str = "",
|
||||
min_rating: int = 0,
|
||||
limit: int = 50,
|
||||
skip_multipart: bool = True,
|
||||
):
|
||||
"""
|
||||
Search for prompts matching the given criteria.
|
||||
|
||||
Args:
|
||||
text: Search text to match against prompt content
|
||||
category: Filter by specific category
|
||||
tags: Comma-separated list of tags to filter by
|
||||
min_rating: Minimum rating (0-5) to include
|
||||
limit: Maximum number of results to return
|
||||
|
||||
Returns:
|
||||
Dict with 'ui' display info and 'result' tuple
|
||||
"""
|
||||
try:
|
||||
# Parse tags from comma-separated string
|
||||
tags_list = None
|
||||
if tags and tags.strip():
|
||||
tags_list = [tag.strip() for tag in tags.split(",") if tag.strip()]
|
||||
|
||||
# Perform search with partial tag matching for discovery
|
||||
results = self.db.search_prompts(
|
||||
text=text.strip() if text and text.strip() else None,
|
||||
category=category.strip() if category and category.strip() else None,
|
||||
tags=tags_list,
|
||||
tag_partial=True,
|
||||
rating_min=min_rating if min_rating > 0 else None,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
# Extract prompt text, collapsing newlines to spaces so downstream
|
||||
# nodes that split on \n (e.g. StringOutputList) treat each DB
|
||||
# entry as a single prompt.
|
||||
prompt_texts = [
|
||||
" ".join(r["text"].split()) for r in results if r.get("text")
|
||||
]
|
||||
|
||||
# Filter out multi-part prompts (Clip_1/Clip_2 markers)
|
||||
if skip_multipart:
|
||||
prompt_texts = [
|
||||
p for p in prompt_texts if not re.search(r"Clip_\d+", p)
|
||||
]
|
||||
|
||||
# Filter out prompts that are only LoRA tags with no actual content.
|
||||
# Uses subtraction instead of repeated groups to avoid backtracking.
|
||||
prompt_texts = [
|
||||
p for p in prompt_texts if re.sub(r"<lora:[^>]+>", "", p).strip()
|
||||
]
|
||||
|
||||
count = len(prompt_texts)
|
||||
|
||||
# Build a preview: numbered list of truncated prompts
|
||||
preview_lines = []
|
||||
for i, p in enumerate(prompt_texts, 1):
|
||||
truncated = p[:120] + "..." if len(p) > 120 else p
|
||||
preview_lines.append(f"[{i}] {truncated}")
|
||||
preview = "\n".join(preview_lines) if preview_lines else "No results found"
|
||||
|
||||
# Must return at least one element — ComfyUI's slice_dict
|
||||
# indexes into OUTPUT_IS_LIST outputs and crashes on empty lists.
|
||||
if not prompt_texts:
|
||||
self.logger.info("Search returned no results")
|
||||
return {
|
||||
"ui": {"text": ["No results found"]},
|
||||
"result": ([""], preview, count),
|
||||
}
|
||||
|
||||
return {
|
||||
"ui": {"text": [f"Found {count} prompts"]},
|
||||
"result": (prompt_texts, preview, count),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Search error: {e}", exc_info=True)
|
||||
return {
|
||||
"ui": {"text": [f"Search error: {e}"]},
|
||||
"result": ([""], f"Error: {e}", 0),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(
|
||||
cls, text="", category="", tags="", min_rating=0, limit=50, skip_multipart=True
|
||||
):
|
||||
"""
|
||||
Always re-execute to get fresh results from the database.
|
||||
|
||||
The database contents may have changed since the last execution,
|
||||
so we always return a unique value to trigger re-execution.
|
||||
"""
|
||||
import time
|
||||
|
||||
return time.time()
|
||||
@@ -3,4 +3,4 @@ ComfyUI_PromptManager Python API modules.
|
||||
|
||||
This package contains the core API components for the web interface and configuration
|
||||
management, including REST endpoints, configuration handling, and server integration.
|
||||
"""
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,972 @@
|
||||
"""REST API module for ComfyUI PromptManager.
|
||||
|
||||
Split into domain-specific mixins for maintainability:
|
||||
- PromptRoutesMixin: prompt CRUD, tags, categories, bulk ops, export
|
||||
- ImageRoutesMixin: gallery, thumbnails, image serving and linking
|
||||
- AdminRoutesMixin: duplicates, stats, settings, diagnostics, maintenance, backup, scan
|
||||
- LoggingRoutesMixin: log management endpoints
|
||||
- AutotagRoutesMixin: auto-tagging model management and tagging
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import functools
|
||||
import gzip as gzip_module
|
||||
import json
|
||||
import ntpath
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
|
||||
from .prompts import PromptRoutesMixin
|
||||
from .images import ImageRoutesMixin
|
||||
from .admin import AdminRoutesMixin
|
||||
from .logging_routes import LoggingRoutesMixin
|
||||
from .autotag_routes import AutotagRoutesMixin
|
||||
from .lora_integration import LoraIntegrationMixin
|
||||
|
||||
try:
|
||||
from ...database.operations import PromptDatabase
|
||||
from ...utils.logging_config import get_logger
|
||||
from ...utils.video_metadata import is_video_path, read_video_metadata
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
)
|
||||
from database.operations import PromptDatabase
|
||||
from utils.logging_config import get_logger
|
||||
from utils.video_metadata import is_video_path, read_video_metadata
|
||||
|
||||
|
||||
def _get_project_root():
|
||||
"""Get the project root directory (3 levels up from py/api/__init__.py)."""
|
||||
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
|
||||
def _public_path(path) -> str:
|
||||
"""Render a server path for an API response without revealing the layout.
|
||||
|
||||
Returns the path relative to the ComfyUI base directory (or to the parent
|
||||
of an opted-in extra gallery root), as a POSIX-style string, or just the
|
||||
basename when the path is not under any of those anchors.
|
||||
"""
|
||||
if not path:
|
||||
return ""
|
||||
from ..config import GalleryConfig
|
||||
|
||||
try:
|
||||
canonical = Path(os.path.normcase(os.path.realpath(str(path))))
|
||||
except (OSError, ValueError):
|
||||
return os.path.basename(str(path))
|
||||
for anchor in GalleryConfig.path_anchors():
|
||||
if canonical.is_relative_to(Path(anchor)):
|
||||
return canonical.relative_to(Path(anchor)).as_posix() or "."
|
||||
return canonical.name
|
||||
|
||||
|
||||
def _public_error(exc: BaseException) -> str:
|
||||
"""Error text safe for API responses: OSError paths reduced to basenames."""
|
||||
if isinstance(exc, OSError):
|
||||
parts = [exc.strerror or type(exc).__name__]
|
||||
if exc.filename:
|
||||
parts.append(os.path.basename(str(exc.filename)))
|
||||
return ": ".join(parts)
|
||||
return str(exc) or type(exc).__name__
|
||||
|
||||
|
||||
# ── Static file serving ────────────────────────────────────────────────
|
||||
# Only these extensions are ever served from web/lib and web/js; the map
|
||||
# doubles as the allow-list.
|
||||
STATIC_MIME_TYPES = {
|
||||
".js": "application/javascript",
|
||||
".css": "text/css",
|
||||
".map": "application/json",
|
||||
".woff": "font/woff",
|
||||
".woff2": "font/woff2",
|
||||
".ttf": "font/ttf",
|
||||
".png": "image/png",
|
||||
".svg": "image/svg+xml",
|
||||
".ico": "image/x-icon",
|
||||
".html": "text/html",
|
||||
}
|
||||
_STATIC_SEPARATORS = re.compile(r"[\\/]+")
|
||||
|
||||
|
||||
def _resolve_static_file(root, filepath):
|
||||
"""Real path of ``filepath`` under ``root``, or None when it must not be served.
|
||||
|
||||
The request path is untrusted. It is rejected when it is empty, carries a
|
||||
``..`` segment, is absolute on any platform (a leading slash or backslash,
|
||||
a drive letter such as ``C:`` or a UNC ``\\\\server\\share`` prefix all
|
||||
re-anchor ``os.path.join`` on Windows), has an extension outside
|
||||
STATIC_MIME_TYPES, or resolves (symlinks included) outside ``root``.
|
||||
"""
|
||||
if not isinstance(filepath, str) or not filepath:
|
||||
return None
|
||||
if (
|
||||
os.path.isabs(filepath)
|
||||
or ntpath.isabs(filepath)
|
||||
or ntpath.splitdrive(filepath)[0]
|
||||
):
|
||||
return None
|
||||
segments = _STATIC_SEPARATORS.split(filepath)
|
||||
if any(segment in ("", ".", "..") for segment in segments):
|
||||
return None
|
||||
if os.path.splitext(filepath)[1].lower() not in STATIC_MIME_TYPES:
|
||||
return None
|
||||
try:
|
||||
root_real = Path(os.path.normcase(os.path.realpath(root)))
|
||||
candidate = Path(
|
||||
os.path.normcase(os.path.realpath(os.path.join(root, *segments)))
|
||||
)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
if candidate == root_real or not candidate.is_relative_to(root_real):
|
||||
return None
|
||||
return str(candidate)
|
||||
|
||||
|
||||
# ── JSON request bodies ────────────────────────────────────────────────
|
||||
JSON_BODY_MAX_BYTES = 1_000_000
|
||||
_JSON_READ_CHUNK = 64 * 1024
|
||||
|
||||
|
||||
def _json_error(message, status):
|
||||
return web.json_response({"success": False, "error": message}, status=status)
|
||||
|
||||
|
||||
async def _read_json_body(request, max_bytes=JSON_BODY_MAX_BYTES):
|
||||
"""Read and parse a JSON object body of at most ``max_bytes``.
|
||||
|
||||
The body is read in chunks so an oversized or chunked upload is refused
|
||||
as soon as it crosses the cap rather than buffered whole.
|
||||
|
||||
Returns:
|
||||
(data, None) on success, or (None, response) where ``response`` is a
|
||||
413 (body too large) or 400 (invalid JSON, or not a JSON object).
|
||||
"""
|
||||
too_large = _json_error(f"Request body too large (max {max_bytes} bytes)", 413)
|
||||
declared = request.content_length
|
||||
if declared is not None and declared > max_bytes:
|
||||
return None, too_large
|
||||
|
||||
chunks = []
|
||||
size = 0
|
||||
async for chunk in request.content.iter_chunked(_JSON_READ_CHUNK):
|
||||
size += len(chunk)
|
||||
if size > max_bytes:
|
||||
return None, too_large
|
||||
chunks.append(chunk)
|
||||
|
||||
try:
|
||||
data = json.loads(b"".join(chunks).decode("utf-8"))
|
||||
except (UnicodeDecodeError, ValueError):
|
||||
return None, _json_error("Request body must be valid JSON", 400)
|
||||
if not isinstance(data, dict):
|
||||
return None, _json_error("Request body must be a JSON object", 400)
|
||||
return data, None
|
||||
|
||||
|
||||
class _LongJobGuard:
|
||||
"""Async context manager handed out by ``_acquire_long_job``."""
|
||||
|
||||
def __init__(self, lock):
|
||||
self._lock = lock
|
||||
|
||||
async def __aenter__(self):
|
||||
# The caller checked ``locked()`` on this same loop iteration, so this
|
||||
# acquire never waits.
|
||||
await self._lock.acquire()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
self._lock.release()
|
||||
return False
|
||||
|
||||
|
||||
# ── Gzip compression middleware ────────────────────────────────────────
|
||||
_GZIP_MIN_SIZE = 1024 # Only compress bodies larger than 1 KB
|
||||
_GZIP_TYPES = frozenset(
|
||||
(
|
||||
"application/json",
|
||||
"text/html",
|
||||
"text/css",
|
||||
"application/javascript",
|
||||
"text/plain",
|
||||
)
|
||||
)
|
||||
_gzip_registered = False
|
||||
|
||||
|
||||
@web.middleware
|
||||
async def _gzip_middleware(request, handler):
|
||||
"""Compress PromptManager responses when the client accepts gzip."""
|
||||
response = await handler(request)
|
||||
|
||||
if not request.path.startswith("/prompt_manager/"):
|
||||
return response
|
||||
|
||||
# Only handle regular Response objects (not StreamResponse/WebSocket)
|
||||
if not isinstance(response, web.Response):
|
||||
return response
|
||||
|
||||
if "gzip" not in request.headers.get("Accept-Encoding", ""):
|
||||
return response
|
||||
|
||||
if "Content-Encoding" in response.headers:
|
||||
return response
|
||||
|
||||
body = response.body
|
||||
if body is None or len(body) < _GZIP_MIN_SIZE:
|
||||
return response
|
||||
|
||||
content_type = response.content_type or ""
|
||||
if not any(ct in content_type for ct in _GZIP_TYPES):
|
||||
return response
|
||||
|
||||
compressed = gzip_module.compress(body, compresslevel=6)
|
||||
if len(compressed) >= len(body):
|
||||
return response
|
||||
|
||||
response.body = compressed
|
||||
response.headers["Content-Encoding"] = "gzip"
|
||||
response.headers["Vary"] = "Accept-Encoding"
|
||||
return response
|
||||
|
||||
|
||||
class PromptManagerAPI(
|
||||
PromptRoutesMixin,
|
||||
ImageRoutesMixin,
|
||||
AdminRoutesMixin,
|
||||
LoggingRoutesMixin,
|
||||
AutotagRoutesMixin,
|
||||
LoraIntegrationMixin,
|
||||
):
|
||||
"""REST API handler for PromptManager operations and web interface.
|
||||
|
||||
This class provides comprehensive REST API endpoints for managing prompts,
|
||||
images, and system operations. It handles database interactions, file
|
||||
operations, image processing, and web UI serving.
|
||||
|
||||
The API is designed to integrate seamlessly with ComfyUI's aiohttp server
|
||||
and provides both JSON API endpoints and static file serving for the
|
||||
web interface.
|
||||
|
||||
Attributes:
|
||||
logger: Configured logger instance for API operations
|
||||
db (PromptDatabase): Database connection and operations handler
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the PromptManager API with database connection and cleanup."""
|
||||
self.logger = get_logger("prompt_manager.api")
|
||||
self.logger.info("Initializing PromptManager API")
|
||||
|
||||
self.db = PromptDatabase()
|
||||
self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir()
|
||||
self._html_cache = {} # Cached HTML file contents keyed by path
|
||||
self._gallery_cache = {} # dict: path_str -> (files, timestamp)
|
||||
self._gallery_cache_ttl = 30 # Cache TTL in seconds
|
||||
# Single-flight locks for long jobs, keyed by job family; see
|
||||
# _acquire_long_job(). Created lazily so no lock binds to an event
|
||||
# loop before the server's loop exists.
|
||||
self._long_jobs: dict = {}
|
||||
# The current output scan (py/api/scan_job.ScanJob), if one was started.
|
||||
self._scan_job = None
|
||||
# The current duplicate scan (also a ScanJob), if one was started.
|
||||
self._duplicate_job = None
|
||||
|
||||
# Run cleanup on initialization to remove any existing duplicates
|
||||
try:
|
||||
removed = self.db.cleanup_duplicates()
|
||||
if removed > 0:
|
||||
self.logger.info(
|
||||
f"Startup cleanup: removed {removed} duplicate prompts"
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Startup cleanup failed: {e}")
|
||||
|
||||
self.logger.info("PromptManager API initialization completed")
|
||||
|
||||
_public_path = staticmethod(_public_path)
|
||||
_public_error = staticmethod(_public_error)
|
||||
_read_json_body = staticmethod(_read_json_body)
|
||||
|
||||
def _acquire_long_job(self, name):
|
||||
"""Claim the single-flight slot for a long job family, or None when busy.
|
||||
|
||||
Usage in a handler::
|
||||
|
||||
job = self._acquire_long_job("duplicates")
|
||||
if job is None:
|
||||
return self._long_job_busy_response("Duplicate scan")
|
||||
async with job:
|
||||
...
|
||||
|
||||
Families are free-form strings ("duplicates", "thumbnails",
|
||||
"autotag", ...); the lock is released when the ``async with`` block
|
||||
exits, including on exceptions. Must be called from the event loop.
|
||||
"""
|
||||
lock = self._long_jobs.setdefault(name, asyncio.Lock())
|
||||
if lock.locked():
|
||||
return None
|
||||
return _LongJobGuard(lock)
|
||||
|
||||
@staticmethod
|
||||
def _long_job_busy_response(label):
|
||||
"""409 envelope for a long job that is already running."""
|
||||
return _json_error(f"{label} already running", 409)
|
||||
|
||||
async def _run_in_executor(self, func, *args, **kwargs):
|
||||
"""Run a blocking function in the default thread pool executor.
|
||||
|
||||
Prevents synchronous database calls, file I/O, and PIL operations
|
||||
from blocking the aiohttp event loop.
|
||||
"""
|
||||
loop = asyncio.get_event_loop()
|
||||
if kwargs:
|
||||
call = functools.partial(func, *args, **kwargs)
|
||||
return await loop.run_in_executor(None, call)
|
||||
return await loop.run_in_executor(None, func, *args)
|
||||
|
||||
def invalidate_gallery_cache(self):
|
||||
"""Invalidate the cached gallery file listing.
|
||||
|
||||
Called by the image monitor when new files are detected.
|
||||
"""
|
||||
self._gallery_cache = {}
|
||||
|
||||
def add_routes(self, routes):
|
||||
"""Register all API routes with the ComfyUI server.
|
||||
|
||||
Args:
|
||||
routes: aiohttp RouteTableDef object from ComfyUI server instance.
|
||||
"""
|
||||
|
||||
# Test route to verify registration works
|
||||
@routes.get("/prompt_manager/test")
|
||||
async def test_route(request):
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"message": "PromptManager API is working!",
|
||||
"timestamp": str(datetime.datetime.now()),
|
||||
}
|
||||
)
|
||||
|
||||
# ── Web UI serving routes ─────────────────────────────────────
|
||||
|
||||
@routes.get("/prompt_manager/web")
|
||||
async def serve_web_ui(request):
|
||||
return self._serve_html_page("index.html", "Web UI", cache=False)
|
||||
|
||||
@routes.get("/prompt_manager/gallery.html")
|
||||
async def serve_gallery_ui(request):
|
||||
return self._serve_html_page("metadata.html", "Gallery")
|
||||
|
||||
@routes.get("/prompt_manager/admin")
|
||||
async def serve_admin_ui(request):
|
||||
return self._serve_html_page("admin.html", "Admin UI")
|
||||
|
||||
@routes.get("/prompt_manager/gallery")
|
||||
async def serve_gallery_admin_ui(request):
|
||||
return self._serve_html_page("gallery.html", "Gallery")
|
||||
|
||||
# ── Static file serving ───────────────────────────────────────
|
||||
|
||||
@routes.get("/prompt_manager/lib/{filepath:.*}")
|
||||
async def serve_lib_static(request):
|
||||
"""Serve static library files (JS, CSS) from web/lib directory."""
|
||||
return self._serve_static_file("lib", request.match_info.get("filepath"))
|
||||
|
||||
@routes.get("/prompt_manager/js/{filepath:.*}")
|
||||
async def serve_js_static(request):
|
||||
"""Serve static JavaScript files from web/js directory."""
|
||||
return self._serve_static_file("js", request.match_info.get("filepath"))
|
||||
|
||||
# ── Register domain-specific routes from mixins ───────────────
|
||||
|
||||
self._register_prompt_routes(routes)
|
||||
self._register_image_routes(routes)
|
||||
self._register_admin_routes(routes)
|
||||
self._register_logging_routes(routes)
|
||||
self._register_autotag_routes(routes)
|
||||
self._register_lora_routes(routes)
|
||||
|
||||
# Register gzip compression middleware (once)
|
||||
global _gzip_registered
|
||||
if not _gzip_registered:
|
||||
try:
|
||||
from ..config import server_instance
|
||||
|
||||
server_instance.app.middlewares.append(_gzip_middleware)
|
||||
_gzip_registered = True
|
||||
self.logger.info("Gzip compression middleware registered")
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not register gzip middleware: {e}")
|
||||
|
||||
self.logger.info("All routes registered with decorator pattern")
|
||||
|
||||
def _serve_static_file(self, subdir, filepath):
|
||||
"""Serve ``web/<subdir>/<filepath>``; 403 outside it, 404 when missing."""
|
||||
root = os.path.join(_get_project_root(), "web", subdir)
|
||||
file_path = _resolve_static_file(root, filepath)
|
||||
if file_path is None:
|
||||
return web.Response(text="Forbidden", status=403)
|
||||
if not os.path.isfile(file_path):
|
||||
return web.Response(text="Not Found", status=404)
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
return web.Response(body=content, content_type=STATIC_MIME_TYPES[ext])
|
||||
|
||||
def _serve_html_page(self, filename, title, cache=True):
|
||||
"""Serve ``web/<filename>``; 404 when missing, 500 when unreadable."""
|
||||
html_path = os.path.join(_get_project_root(), "web", filename)
|
||||
try:
|
||||
if cache and html_path in self._html_cache:
|
||||
content = self._html_cache[html_path]
|
||||
elif os.path.exists(html_path):
|
||||
with open(html_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
if cache:
|
||||
self._html_cache[html_path] = content
|
||||
else:
|
||||
return web.Response(
|
||||
text=f"<h1>{title} not found</h1>"
|
||||
"<p>HTML file not located at expected path.</p>",
|
||||
content_type="text/html",
|
||||
status=404,
|
||||
)
|
||||
return web.Response(text=content, content_type="text/html", charset="utf-8")
|
||||
except Exception:
|
||||
self.logger.exception(f"Failed to load {title}")
|
||||
return web.Response(
|
||||
text=f"<h1>Error</h1><p>Failed to load {title}. "
|
||||
"Check server logs for details.</p>",
|
||||
content_type="text/html",
|
||||
status=500,
|
||||
)
|
||||
|
||||
# ── Shared utilities used by multiple mixins ──────────────────────
|
||||
|
||||
def _enrich_images(self, images):
|
||||
"""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()
|
||||
|
||||
for image in images:
|
||||
image_path_str = image.get("image_path", "")
|
||||
if not image_path_str:
|
||||
continue
|
||||
|
||||
img_path = Path(image_path_str)
|
||||
|
||||
# Set fallback url via image ID
|
||||
if image.get("id"):
|
||||
image["url"] = f"/prompt_manager/images/{image['id']}/file"
|
||||
|
||||
# Try each root to compute relative path and thumbnail URL
|
||||
for output_path in output_dirs:
|
||||
try:
|
||||
rel_path = img_path.resolve().relative_to(output_path.resolve())
|
||||
image["relative_path"] = str(rel_path)
|
||||
image["url"] = (
|
||||
"/prompt_manager/images/serve/"
|
||||
f"{url_quote(rel_path.as_posix(), safe='/')}"
|
||||
)
|
||||
|
||||
# Check for thumbnail
|
||||
rel_no_ext = rel_path.with_suffix("")
|
||||
thumb_rel = (
|
||||
f"thumbnails/{rel_no_ext.as_posix()}_thumb{rel_path.suffix}"
|
||||
)
|
||||
from .prompts import thumbnail_url_for
|
||||
|
||||
thumbnail_url = thumbnail_url_for(
|
||||
thumb_rel, output_path / thumb_rel
|
||||
)
|
||||
if thumbnail_url:
|
||||
image["thumbnail_url"] = thumbnail_url
|
||||
break # Found matching root, stop searching
|
||||
except (ValueError, RuntimeError):
|
||||
continue # Try next root
|
||||
|
||||
return images
|
||||
|
||||
def _enrich_prompt_images(self, prompts):
|
||||
"""Add url and thumbnail_url to each image in prompt results."""
|
||||
for prompt in prompts:
|
||||
self._enrich_images(prompt.get("images", []))
|
||||
return prompts
|
||||
|
||||
def _clean_nan_recursive(self, obj):
|
||||
"""Recursively clean NaN values from nested data structures."""
|
||||
if isinstance(obj, dict):
|
||||
return {key: self._clean_nan_recursive(value) for key, value in obj.items()}
|
||||
elif isinstance(obj, list):
|
||||
return [self._clean_nan_recursive(item) for item in obj]
|
||||
elif isinstance(obj, float) and str(obj) == "nan":
|
||||
return None
|
||||
else:
|
||||
return obj
|
||||
|
||||
def _find_comfyui_output_dir(self):
|
||||
"""Locate the ComfyUI output directory using multiple detection strategies.
|
||||
|
||||
Checks in order:
|
||||
1. User-configured directory (GalleryConfig.MONITORING_DIRECTORIES)
|
||||
2. Upward filesystem search for ComfyUI markers
|
||||
3. Common installation path patterns
|
||||
4. Common ComfyUI installation locations
|
||||
|
||||
Results are cached after first successful lookup. The cache is
|
||||
invalidated when the user changes the configured directory via
|
||||
save_settings().
|
||||
|
||||
Returns:
|
||||
str or None: Absolute path to ComfyUI output directory, or None
|
||||
if no valid directory is found
|
||||
"""
|
||||
if self._cached_output_dir is not None:
|
||||
return self._cached_output_dir
|
||||
|
||||
# Method 0: Check user-configured directory first
|
||||
try:
|
||||
from ..config import GalleryConfig
|
||||
|
||||
if GalleryConfig.MONITORING_DIRECTORIES:
|
||||
configured = GalleryConfig.MONITORING_DIRECTORIES[0]
|
||||
# Same check as the settings endpoint: a root hand-edited
|
||||
# into config.json must not open a directory outside ComfyUI
|
||||
ok, reason = GalleryConfig.validate_gallery_root(configured)
|
||||
if ok:
|
||||
configured_dir = Path(configured).resolve()
|
||||
self.logger.info(
|
||||
f"Using configured monitoring directory: {configured_dir}"
|
||||
)
|
||||
self._cached_output_dir = str(configured_dir)
|
||||
return self._cached_output_dir
|
||||
self.logger.warning(f"Ignoring configured directory: {reason}")
|
||||
except ImportError:
|
||||
self.logger.debug(
|
||||
"GalleryConfig not available, skipping configured directory check"
|
||||
)
|
||||
|
||||
current_file = Path(__file__).resolve()
|
||||
self.logger.debug(f"Starting ComfyUI output search from: {current_file}")
|
||||
|
||||
# Method 1: Search upward from current file location
|
||||
current_dir = current_file.parent
|
||||
max_depth = 10 # Prevent infinite loops
|
||||
|
||||
for i in range(max_depth):
|
||||
# Check if current directory contains ComfyUI markers
|
||||
comfyui_markers = ["main.py", "nodes.py", "server.py"]
|
||||
if any((current_dir / marker).exists() for marker in comfyui_markers):
|
||||
output_dir = current_dir / "output"
|
||||
if output_dir.exists() and output_dir.is_dir():
|
||||
self.logger.debug(
|
||||
"Found ComfyUI output directory via upward search: "
|
||||
f"{output_dir}"
|
||||
)
|
||||
self._cached_output_dir = str(output_dir)
|
||||
return self._cached_output_dir
|
||||
|
||||
# Move up one directory
|
||||
parent = current_dir.parent
|
||||
if parent == current_dir: # Reached filesystem root
|
||||
break
|
||||
current_dir = parent
|
||||
|
||||
# Method 2: Try common installation patterns relative to this file
|
||||
# File is at: custom_nodes/ComfyUI_PromptManager/py/api/__init__.py
|
||||
base_dir = current_file.parent # .../py/api/
|
||||
possible_paths = [
|
||||
base_dir.parent.parent.parent.parent / "output", # ../../../../output
|
||||
base_dir.parent.parent.parent / "output", # ../../../output
|
||||
base_dir.parent.parent / "output", # ../../output
|
||||
base_dir.parent / "output", # ../output
|
||||
]
|
||||
|
||||
# Method 3: Add common ComfyUI installation locations
|
||||
common_locations = [
|
||||
Path.home() / "ComfyUI" / "output",
|
||||
Path.cwd() / "output",
|
||||
Path.cwd() / ".." / "output",
|
||||
Path.cwd() / ".." / ".." / "output",
|
||||
]
|
||||
|
||||
all_paths = possible_paths + common_locations
|
||||
|
||||
for path in all_paths:
|
||||
try:
|
||||
abs_path = path.resolve()
|
||||
if abs_path.exists() and abs_path.is_dir():
|
||||
self.logger.debug(f"Found ComfyUI output directory: {abs_path}")
|
||||
self._cached_output_dir = str(abs_path)
|
||||
return self._cached_output_dir
|
||||
except (OSError, RuntimeError):
|
||||
continue # Skip invalid paths
|
||||
|
||||
self.logger.warning("ComfyUI output directory not found. Searched paths:")
|
||||
for path in all_paths:
|
||||
try:
|
||||
self.logger.warning(f" - {path.resolve()} (exists: {path.exists()})")
|
||||
except (OSError, RuntimeError):
|
||||
self.logger.warning(f" - {path} (invalid path)")
|
||||
|
||||
return None
|
||||
|
||||
def _get_all_output_dirs(self):
|
||||
"""Get the directories images may be served from.
|
||||
|
||||
Configured gallery roots are filtered through
|
||||
``GalleryConfig.validate_gallery_root`` so a root that was hand-edited
|
||||
into config.json (or persisted before validation existed) can never
|
||||
expose files outside ComfyUI's directories.
|
||||
|
||||
Returns:
|
||||
List[Path]: Validated directory paths, possibly empty.
|
||||
"""
|
||||
from ..config import GalleryConfig
|
||||
|
||||
configured = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
output_dirs = []
|
||||
seen = set()
|
||||
for d in configured:
|
||||
ok, reason = GalleryConfig.validate_gallery_root(d)
|
||||
if not ok:
|
||||
self.logger.warning(f"Ignoring configured gallery root: {reason}")
|
||||
continue
|
||||
real = os.path.realpath(d)
|
||||
key = os.path.normcase(real) # one entry per directory, however spelled
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
output_dirs.append(Path(real))
|
||||
|
||||
if output_dirs:
|
||||
return output_dirs
|
||||
|
||||
if configured:
|
||||
# Every configured root was rejected: never let auto-detection
|
||||
# resurrect it, serve ComfyUI's own output directory instead.
|
||||
return self._comfyui_output_dir_list()
|
||||
|
||||
fallback = self._find_comfyui_output_dir()
|
||||
return [Path(fallback)] if fallback else []
|
||||
|
||||
def _comfyui_output_dir_list(self):
|
||||
"""ComfyUI's output directory as a one-element list, or [] if unknown."""
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
except (ImportError, AttributeError, OSError):
|
||||
return []
|
||||
if isinstance(output_dir, str) and os.path.isdir(output_dir):
|
||||
return [Path(os.path.realpath(output_dir))]
|
||||
return []
|
||||
|
||||
def _extract_comfyui_metadata(self, image_path):
|
||||
"""Embedded ComfyUI metadata of an image (PNG text chunks) or a video.
|
||||
|
||||
Videos go through ffprobe (see utils/video_metadata). Unreadable files
|
||||
are a warning, not an error: a scan over a large output tree meets
|
||||
truncated and foreign files, and it carries on past them.
|
||||
"""
|
||||
if is_video_path(image_path):
|
||||
return read_video_metadata(image_path)
|
||||
try:
|
||||
with Image.open(image_path) as img:
|
||||
return dict(getattr(img, "text", None) or {})
|
||||
except Exception as e:
|
||||
self.logger.warning(
|
||||
f"Unreadable image {os.path.basename(str(image_path))}: {e}"
|
||||
)
|
||||
return {}
|
||||
|
||||
def _parse_comfyui_prompt(self, metadata):
|
||||
"""Parse ComfyUI and A1111 prompt data from metadata."""
|
||||
result = {
|
||||
"prompt": None,
|
||||
"workflow": None,
|
||||
"parameters": {},
|
||||
"positive_prompt": None,
|
||||
"negative_prompt": None,
|
||||
}
|
||||
|
||||
# Check for A1111/Forge style parameters first.
|
||||
# The chunk is one block: a possibly multi-line positive prompt, then a
|
||||
# "Negative prompt:" section, then a trailing "Steps: ..." settings line.
|
||||
# Taking lines[0] as the positive prompt truncates any multi-line prompt
|
||||
# to one line, so split on the settings line instead.
|
||||
if "parameters" in metadata:
|
||||
import re as _re
|
||||
|
||||
params = metadata["parameters"]
|
||||
body = params
|
||||
settings = ""
|
||||
_sm = _re.search(r"\nSteps:\s", body)
|
||||
if _sm:
|
||||
settings = body[_sm.start() + 1 :].strip()
|
||||
body = body[: _sm.start()]
|
||||
negative = None
|
||||
_nm = _re.search(r"(?im)^\s*negative prompt:\s*", body)
|
||||
if _nm:
|
||||
negative = body[_nm.end() :].strip()
|
||||
body = body[: _nm.start()]
|
||||
elif body.lower().lstrip().startswith("negative prompt:"):
|
||||
negative = body.lower().lstrip()[len("negative prompt:") :].strip()
|
||||
body = ""
|
||||
result["positive_prompt"] = body.strip() or None
|
||||
if negative:
|
||||
result["negative_prompt"] = negative
|
||||
# Store raw parameters too
|
||||
result["parameters"]["parameters"] = params
|
||||
if settings:
|
||||
for _f in ("Steps", "Sampler", "Schedule type", "CFG scale", "Seed"):
|
||||
_fm = _re.search(
|
||||
r"(?i)(?:^|,\s*)" + _re.escape(_f) + r":\s*([^,]+)", settings
|
||||
)
|
||||
if _fm:
|
||||
result["parameters"][_f.lower().replace(" ", "_")] = _fm.group(
|
||||
1
|
||||
).strip()
|
||||
|
||||
# If no A1111 format found, proceed with ComfyUI parsing
|
||||
if result["positive_prompt"] is None:
|
||||
# Check for direct prompt field
|
||||
if "prompt" in metadata:
|
||||
try:
|
||||
prompt_data = json.loads(metadata["prompt"])
|
||||
result["prompt"] = prompt_data
|
||||
except json.JSONDecodeError:
|
||||
result["prompt"] = metadata["prompt"]
|
||||
|
||||
# Check for workflow
|
||||
if "workflow" in metadata:
|
||||
try:
|
||||
workflow_data = json.loads(metadata["workflow"])
|
||||
result["workflow"] = workflow_data
|
||||
except json.JSONDecodeError:
|
||||
result["workflow"] = metadata["workflow"]
|
||||
|
||||
# Check for other common ComfyUI fields
|
||||
common_fields = [
|
||||
"positive",
|
||||
"negative",
|
||||
"steps",
|
||||
"cfg",
|
||||
"sampler",
|
||||
"scheduler",
|
||||
"seed",
|
||||
]
|
||||
for field in common_fields:
|
||||
if field in metadata:
|
||||
try:
|
||||
result["parameters"][field] = json.loads(metadata[field])
|
||||
except json.JSONDecodeError:
|
||||
result["parameters"][field] = metadata[field]
|
||||
|
||||
return result
|
||||
|
||||
def _extract_readable_prompt(self, parsed_data):
|
||||
"""Extract human-readable prompt text from ComfyUI/A1111 data."""
|
||||
|
||||
def safe_to_string(value):
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
elif isinstance(value, list):
|
||||
return " ".join(str(item) for item in value if item)
|
||||
elif value is not None:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
# First check if we already extracted a positive prompt (A1111 format)
|
||||
if parsed_data.get("positive_prompt"):
|
||||
return parsed_data["positive_prompt"]
|
||||
|
||||
# Check if prompt is already a string
|
||||
if isinstance(parsed_data.get("prompt"), str):
|
||||
return parsed_data["prompt"]
|
||||
|
||||
# Check if prompt is a simple value that can be converted
|
||||
if parsed_data.get("prompt") and not isinstance(
|
||||
parsed_data.get("prompt"), dict
|
||||
):
|
||||
return safe_to_string(parsed_data["prompt"])
|
||||
|
||||
prompt_data = parsed_data.get("prompt")
|
||||
if isinstance(prompt_data, dict):
|
||||
# Use the enhanced logic from parse-metadata.py
|
||||
positive_prompt = self._extract_positive_prompt_from_comfyui_data(
|
||||
prompt_data
|
||||
)
|
||||
if positive_prompt:
|
||||
return positive_prompt
|
||||
|
||||
# Check workflow data if available
|
||||
workflow_data = parsed_data.get("workflow")
|
||||
if isinstance(workflow_data, dict):
|
||||
positive_prompt = self._extract_positive_prompt_from_comfyui_data(
|
||||
workflow_data
|
||||
)
|
||||
if positive_prompt:
|
||||
return positive_prompt
|
||||
|
||||
# Check parameters for positive prompt
|
||||
if parsed_data.get("parameters", {}).get("positive"):
|
||||
return safe_to_string(parsed_data["parameters"]["positive"])
|
||||
|
||||
return None
|
||||
|
||||
def _get_node_inputs(self, node):
|
||||
"""Safely get inputs from a node, handling both dict and list formats.
|
||||
|
||||
Old format: inputs is a dict with direct key-value pairs
|
||||
inputs = {"text": "my prompt", "seed": 123}
|
||||
|
||||
New format: inputs is a list of connection objects
|
||||
inputs = [
|
||||
{"name": "text", "type": "STRING", "link": null,
|
||||
"widget": {"name": "text"}},
|
||||
{"name": "clip", "type": "CLIP", "link": 11}
|
||||
]
|
||||
"""
|
||||
if not isinstance(node, dict):
|
||||
return {}
|
||||
|
||||
inputs = node.get("inputs", {})
|
||||
|
||||
# If inputs is already a dict, return it
|
||||
if isinstance(inputs, dict):
|
||||
return inputs
|
||||
|
||||
# If inputs is a list, convert to dict format
|
||||
if isinstance(inputs, list):
|
||||
inputs_dict = {}
|
||||
for input_item in inputs:
|
||||
if isinstance(input_item, dict) and "name" in input_item:
|
||||
name = input_item["name"]
|
||||
inputs_dict[name] = input_item
|
||||
return inputs_dict
|
||||
|
||||
return {}
|
||||
|
||||
def _find_text_in_node(self, node):
|
||||
"""Try to find text content in a node using various strategies.
|
||||
|
||||
Handles both old and new workflow formats.
|
||||
"""
|
||||
if not isinstance(node, dict):
|
||||
return None
|
||||
|
||||
# Strategy 1: Check normalized inputs for 'text' field
|
||||
inputs = self._get_node_inputs(node)
|
||||
if "text" in inputs and isinstance(inputs["text"], str):
|
||||
return inputs["text"]
|
||||
|
||||
# Strategy 2: For text encoder nodes, check widgets_values
|
||||
class_type = node.get("class_type", node.get("type", ""))
|
||||
text_encoder_types = [
|
||||
"CLIPTextEncode",
|
||||
"CLIPTextEncodeSDXL",
|
||||
"CLIPTextEncodeSDXLRefiner",
|
||||
"CLIPTextEncodeFlux",
|
||||
"PromptManager",
|
||||
"PromptManagerText",
|
||||
"BNK_CLIPTextEncoder",
|
||||
"Text Encoder",
|
||||
"CLIP Text Encode",
|
||||
]
|
||||
|
||||
if any(
|
||||
encoder_type.lower() in class_type.lower()
|
||||
for encoder_type in text_encoder_types
|
||||
):
|
||||
widgets_values = node.get("widgets_values", [])
|
||||
if widgets_values and len(widgets_values) > 0:
|
||||
if isinstance(widgets_values[0], str) and widgets_values[0].strip():
|
||||
return widgets_values[0]
|
||||
|
||||
return None
|
||||
|
||||
def _extract_positive_prompt_from_comfyui_data(self, data):
|
||||
"""Extract positive prompt from ComfyUI data.
|
||||
|
||||
Handles both old and new formats.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
|
||||
# Build nodes dictionary
|
||||
nodes_by_id = {}
|
||||
if "nodes" in data:
|
||||
# Handle nodes array format
|
||||
for node in data["nodes"]:
|
||||
if isinstance(node, dict):
|
||||
nid = node.get("id")
|
||||
if nid is not None:
|
||||
nodes_by_id[nid] = node
|
||||
else:
|
||||
# Handle flat dictionary format (node_id -> node_data)
|
||||
for nid_str, node in data.items():
|
||||
try:
|
||||
nid = int(nid_str)
|
||||
except (ValueError, TypeError):
|
||||
nid = nid_str
|
||||
if isinstance(node, dict):
|
||||
if "id" in node:
|
||||
nid = node["id"]
|
||||
nodes_by_id[nid] = node
|
||||
|
||||
if not nodes_by_id:
|
||||
return None
|
||||
|
||||
# First, try to find positive/negative connection pattern
|
||||
pos_id = None
|
||||
for node in nodes_by_id.values():
|
||||
if isinstance(node, dict):
|
||||
inputs = self._get_node_inputs(node)
|
||||
if "positive" in inputs and "negative" in inputs:
|
||||
try:
|
||||
pos_input = inputs["positive"]
|
||||
if isinstance(pos_input, list) and len(pos_input) > 0:
|
||||
pos_id = int(pos_input[0])
|
||||
break
|
||||
except (ValueError, TypeError, IndexError):
|
||||
continue
|
||||
|
||||
# Get text from the positive node
|
||||
if pos_id is not None and pos_id in nodes_by_id:
|
||||
text_val = self._find_text_in_node(nodes_by_id[pos_id])
|
||||
if text_val:
|
||||
return text_val
|
||||
|
||||
# Fallback: find any text encoder node with text content
|
||||
text_nodes = []
|
||||
for node in nodes_by_id.values():
|
||||
if isinstance(node, dict):
|
||||
text_val = self._find_text_in_node(node)
|
||||
if text_val:
|
||||
# Try to determine if this is positive or negative
|
||||
node_title = node.get("title", "").lower()
|
||||
if "neg" not in node_title and "negative" not in node_title:
|
||||
# Prioritize non-negative prompts
|
||||
text_nodes.insert(0, text_val)
|
||||
else:
|
||||
text_nodes.append(text_val)
|
||||
|
||||
# Return the first positive-looking prompt
|
||||
if text_nodes:
|
||||
return text_nodes[0]
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,746 @@
|
||||
"""AutoTag API routes for PromptManager."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from .prompts import thumbnail_url_for
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
|
||||
_BUSY_ERROR = "Auto-tagging is already running"
|
||||
|
||||
# Only these may be handed to the tagging engine or listed by the scanner.
|
||||
_MEDIA_EXTENSIONS = frozenset(
|
||||
{".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"}
|
||||
)
|
||||
# Bounds for scan_output_dir: a runaway output tree must not pin the server.
|
||||
_SCAN_MAX_FILES = 50_000
|
||||
_SCAN_MAX_DEPTH = 12
|
||||
|
||||
# One batch run at a time; a second autotag/start gets 409 while this is held.
|
||||
_autotag_batch_lock = threading.Lock()
|
||||
|
||||
|
||||
def _file_exists(path) -> bool:
|
||||
try:
|
||||
return Path(str(path)).exists()
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def path_is_within(candidate, root) -> bool:
|
||||
"""True when ``candidate`` is strictly inside ``root`` on the real filesystem.
|
||||
|
||||
Both sides go through ``os.path.realpath`` (symlinks, ``..``) and
|
||||
``os.path.normcase`` (Windows case/separator folding) before the
|
||||
``Path.is_relative_to`` check, so ``/out2`` is never treated as inside
|
||||
``/out`` and a symlink inside ``root`` pointing elsewhere is rejected.
|
||||
"""
|
||||
try:
|
||||
child = Path(os.path.normcase(os.path.realpath(str(candidate))))
|
||||
parent = Path(os.path.normcase(os.path.realpath(str(root))))
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
return False
|
||||
return child != parent and child.is_relative_to(parent)
|
||||
|
||||
|
||||
class AutotagRoutesMixin:
|
||||
"""Mixin providing auto-tagging API endpoints."""
|
||||
|
||||
def _register_autotag_routes(self, routes):
|
||||
@routes.get("/prompt_manager/autotag/models")
|
||||
async def get_autotag_models_route(request):
|
||||
return await self.get_autotag_models(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/download/{model_type}")
|
||||
async def download_autotag_model_route(request):
|
||||
return await self.download_autotag_model(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/start")
|
||||
async def start_autotag_route(request):
|
||||
return await self.start_autotag(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/single")
|
||||
async def autotag_single_route(request):
|
||||
return await self.autotag_single(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/apply")
|
||||
async def apply_autotag_route(request):
|
||||
return await self.apply_autotag(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/unload")
|
||||
async def unload_autotag_model_route(request):
|
||||
return await self.unload_autotag_model(request)
|
||||
|
||||
@routes.get("/prompt_manager/scan_output_dir")
|
||||
async def scan_output_dir_route(request):
|
||||
return await self.scan_output_dir(request)
|
||||
|
||||
async def get_autotag_models(self, request):
|
||||
"""Get status of available AutoTag models."""
|
||||
try:
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
service = get_autotag_service()
|
||||
models_status = service.get_models_status()
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"models": models_status,
|
||||
"default_prompt": service.default_prompt,
|
||||
"model_loaded": service.is_model_loaded(),
|
||||
"loaded_model_type": service.get_loaded_model_type(),
|
||||
"wd14_general_threshold": service.wd14_general_threshold,
|
||||
"wd14_character_threshold": service.wd14_character_threshold,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Get autotag models error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
|
||||
async def _stream_sse(self, request, events):
|
||||
"""Write an async iterator of dict events as server-sent events.
|
||||
|
||||
A client disconnect raises ConnectionResetError from ``write``; the
|
||||
event generator is then closed so its cleanup runs before re-raising
|
||||
(aiohttp logs premature disconnects at debug level).
|
||||
"""
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
reason="OK",
|
||||
headers={
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
},
|
||||
)
|
||||
await response.prepare(request)
|
||||
try:
|
||||
async for event in events:
|
||||
await response.write(f"data: {json.dumps(event)}\n\n".encode("utf-8"))
|
||||
except (ConnectionResetError, asyncio.CancelledError):
|
||||
self.logger.info("Client disconnected; stopping autotag stream")
|
||||
await events.aclose()
|
||||
raise
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
async def download_autotag_model(self, request):
|
||||
"""Download an AutoTag model with streaming progress (POST only)."""
|
||||
model_type = request.match_info.get("model_type")
|
||||
|
||||
async def events():
|
||||
try:
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
service = get_autotag_service()
|
||||
if model_type not in service.models_config:
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": f"Invalid model type: {model_type}",
|
||||
}
|
||||
return
|
||||
|
||||
yield {
|
||||
"type": "progress",
|
||||
"progress": 0,
|
||||
"status": "Starting download...",
|
||||
}
|
||||
|
||||
def progress_callback(status: str, progress: float):
|
||||
pass
|
||||
|
||||
success = await self._run_in_executor(
|
||||
service.download_model, model_type, progress_callback
|
||||
)
|
||||
if success:
|
||||
yield {
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
"status": "Download complete",
|
||||
}
|
||||
else:
|
||||
yield {"type": "error", "message": "Download failed"}
|
||||
except Exception:
|
||||
self.logger.exception("Download model error")
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": _INTERNAL_ERROR,
|
||||
}
|
||||
|
||||
return await self._stream_sse(request, events())
|
||||
|
||||
async def _read_start_params(self, request) -> Optional[dict]:
|
||||
"""Merge query-string and optional JSON-body parameters for autotag/start.
|
||||
|
||||
Returns None when a body is present but is not a JSON object.
|
||||
"""
|
||||
params = dict(request.query)
|
||||
if request.can_read_body:
|
||||
try:
|
||||
body = await request.json()
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
params.update(body)
|
||||
return params
|
||||
|
||||
@staticmethod
|
||||
def _real_tags(tags) -> list:
|
||||
"""Normalise a tag field to a list without bookkeeping markers."""
|
||||
if isinstance(tags, str):
|
||||
tags = [t.strip() for t in tags.split(",") if t.strip()]
|
||||
return [
|
||||
t
|
||||
for t in (tags or [])
|
||||
if t != "auto-scanned"
|
||||
and not t.startswith("prepend:")
|
||||
and not t.startswith("append:")
|
||||
]
|
||||
|
||||
async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool:
|
||||
"""Merge ``tags`` into the prompt; True when something new was added."""
|
||||
existing_prompt = await self._run_in_executor(
|
||||
self.db.get_prompt_by_id, prompt_id
|
||||
)
|
||||
if not existing_prompt:
|
||||
return False
|
||||
existing_tags = existing_prompt.get("tags", [])
|
||||
if isinstance(existing_tags, str):
|
||||
existing_tags = [t.strip() for t in existing_tags.split(",") if t.strip()]
|
||||
new_tags = [t for t in tags if t not in existing_tags]
|
||||
if not new_tags:
|
||||
return False
|
||||
await self._run_in_executor(
|
||||
self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags
|
||||
)
|
||||
return True
|
||||
|
||||
async def _autotag_batch_events(self, service, opts: dict):
|
||||
"""Yield SSE progress dicts while tagging every image linked to a prompt."""
|
||||
import time as _time
|
||||
|
||||
model_type = opts["model_type"]
|
||||
keep_in_memory = opts["keep_in_memory"]
|
||||
status = service.get_models_status()
|
||||
if not status.get(model_type, {}).get("downloaded"):
|
||||
yield {"type": "error", "message": f"Model {model_type} not downloaded"}
|
||||
return
|
||||
|
||||
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
|
||||
try:
|
||||
await self._run_in_executor(service.load_model, model_type, True)
|
||||
except Exception:
|
||||
self.logger.exception("Failed to load model")
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": "Failed to load model. Check server logs for details.",
|
||||
}
|
||||
return
|
||||
|
||||
if opts["custom_prompt"]:
|
||||
service.custom_prompt = opts["custom_prompt"]
|
||||
|
||||
yield {
|
||||
"type": "progress",
|
||||
"progress": 5,
|
||||
"status": "Model loaded. Fetching all images from database...",
|
||||
}
|
||||
images = await self._run_in_executor(self.db.get_all_images)
|
||||
total = len(images)
|
||||
if total == 0:
|
||||
yield {
|
||||
"type": "complete",
|
||||
"processed": 0,
|
||||
"tagged": 0,
|
||||
"skipped": 0,
|
||||
"errors": 0,
|
||||
"skipped_outside_roots": 0,
|
||||
"status": "No images with linked prompts found",
|
||||
}
|
||||
service.unload_model()
|
||||
return
|
||||
|
||||
yield {
|
||||
"type": "progress",
|
||||
"progress": 10,
|
||||
"status": f"Found {total} images. Processing...",
|
||||
}
|
||||
|
||||
counts = {
|
||||
"processed": 0,
|
||||
"tagged": 0,
|
||||
"skipped": 0,
|
||||
"errors": 0,
|
||||
"skipped_outside_roots": 0,
|
||||
}
|
||||
tagged_prompt_ids = set()
|
||||
last_update = _time.monotonic()
|
||||
finished = False
|
||||
|
||||
def progress_event(i, message):
|
||||
return {
|
||||
"type": "progress",
|
||||
"progress": 10 + int((i + 1) / total * 85),
|
||||
"status": message,
|
||||
"processed": counts["processed"],
|
||||
"tagged": counts["tagged"],
|
||||
"skipped": counts["skipped"],
|
||||
"skipped_outside_roots": counts["skipped_outside_roots"],
|
||||
}
|
||||
|
||||
try:
|
||||
for i, image_data in enumerate(images):
|
||||
image_path = image_data.get("image_path")
|
||||
prompt_id = image_data.get("prompt_id")
|
||||
skip_reason = None
|
||||
allowed_path = None
|
||||
if not image_path or not prompt_id:
|
||||
skip_reason = "unlinked"
|
||||
elif prompt_id in tagged_prompt_ids:
|
||||
skip_reason = "prompt already processed"
|
||||
else:
|
||||
# Rows can point anywhere (restored DB, hand edits): only
|
||||
# media files inside the served output roots are tagged.
|
||||
allowed_path = self._resolve_allowed_image_path(str(image_path))
|
||||
if allowed_path is None:
|
||||
if _file_exists(image_path):
|
||||
skip_reason = "outside allowed directories"
|
||||
counts["skipped_outside_roots"] += 1
|
||||
else:
|
||||
skip_reason = "file missing"
|
||||
elif opts["skip_tagged"] and self._real_tags(
|
||||
image_data.get("prompt_tags", [])
|
||||
):
|
||||
tagged_prompt_ids.add(prompt_id)
|
||||
skip_reason = "already tagged"
|
||||
|
||||
if skip_reason is not None:
|
||||
counts["skipped"] += 1
|
||||
now = _time.monotonic()
|
||||
if (now - last_update) >= 0.5 or i == total - 1:
|
||||
yield progress_event(
|
||||
i, f"Skipping {i + 1}/{total} ({skip_reason})..."
|
||||
)
|
||||
await asyncio.sleep(0.01)
|
||||
last_update = now
|
||||
continue
|
||||
|
||||
try:
|
||||
tags = await self._run_in_executor(
|
||||
service.generate_tags,
|
||||
allowed_path,
|
||||
general_threshold=opts["general_threshold"],
|
||||
character_threshold=opts["character_threshold"],
|
||||
)
|
||||
counts["processed"] += 1
|
||||
if tags and await self._apply_tags_to_prompt(prompt_id, tags):
|
||||
counts["tagged"] += 1
|
||||
else:
|
||||
counts["skipped"] += 1
|
||||
tagged_prompt_ids.add(prompt_id)
|
||||
except Exception as img_err:
|
||||
self.logger.error(
|
||||
f"Error processing {os.path.basename(str(image_path))}: "
|
||||
f"{img_err}"
|
||||
)
|
||||
counts["errors"] += 1
|
||||
counts["processed"] += 1
|
||||
|
||||
yield progress_event(i, f"Processing {i + 1}/{total}...")
|
||||
await asyncio.sleep(0.01)
|
||||
last_update = _time.monotonic()
|
||||
|
||||
finished = True
|
||||
if not keep_in_memory:
|
||||
service.unload_model()
|
||||
yield {
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
**counts,
|
||||
"status": "Complete",
|
||||
"model_status": (
|
||||
"Model kept in memory" if keep_in_memory else "Model unloaded"
|
||||
),
|
||||
"model_loaded": keep_in_memory,
|
||||
}
|
||||
finally:
|
||||
if not finished and not keep_in_memory:
|
||||
service.unload_model()
|
||||
|
||||
async def start_autotag(self, request):
|
||||
"""Start batch auto-tagging with streaming progress (POST only)."""
|
||||
params = await self._read_start_params(request)
|
||||
if params is None:
|
||||
return self._bad_request("Request body must be a JSON object")
|
||||
|
||||
opts = {
|
||||
"model_type": params.get("model_type", "gguf"),
|
||||
"custom_prompt": params.get("prompt", ""),
|
||||
"skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true",
|
||||
"keep_in_memory": str(params.get("keep_in_memory", "true")).lower()
|
||||
== "true",
|
||||
"general_threshold": params.get("general_threshold"),
|
||||
"character_threshold": params.get("character_threshold"),
|
||||
}
|
||||
try:
|
||||
for key in ("general_threshold", "character_threshold"):
|
||||
if opts[key] is not None:
|
||||
opts[key] = float(opts[key])
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
)
|
||||
|
||||
async def events():
|
||||
try:
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
async for event in self._autotag_batch_events(
|
||||
get_autotag_service(), opts
|
||||
):
|
||||
yield event
|
||||
except Exception:
|
||||
self.logger.exception("AutoTag error")
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": _INTERNAL_ERROR,
|
||||
}
|
||||
|
||||
if not _autotag_batch_lock.acquire(blocking=False):
|
||||
return web.json_response(
|
||||
{"success": False, "error": _BUSY_ERROR}, status=409
|
||||
)
|
||||
try:
|
||||
return await self._stream_sse(request, events())
|
||||
finally:
|
||||
_autotag_batch_lock.release()
|
||||
|
||||
def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]:
|
||||
"""Return the real path when it is a media file inside an output directory."""
|
||||
if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS:
|
||||
return None
|
||||
for output_dir in self._get_all_output_dirs():
|
||||
if path_is_within(raw_path, output_dir):
|
||||
real = os.path.realpath(raw_path)
|
||||
return real if os.path.isfile(real) else None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _forbidden(raw_path: str):
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": (
|
||||
f"{os.path.basename(str(raw_path).rstrip('/' + os.sep))} is not "
|
||||
"an image inside the configured output directories"
|
||||
),
|
||||
},
|
||||
status=403,
|
||||
)
|
||||
|
||||
async def _select_autotag_target(self, data: dict):
|
||||
"""Pick the image to tag from an autotag/single body.
|
||||
|
||||
Returns ``(image_path, prompt_id, None)`` on success or
|
||||
``(None, None, error_response)`` when the request must be rejected.
|
||||
"""
|
||||
image_id = data.get("image_id")
|
||||
if image_id is not None:
|
||||
if isinstance(image_id, bool) or not isinstance(image_id, (int, str)):
|
||||
return None, None, self._bad_request("image_id must be an integer")
|
||||
try:
|
||||
image_id = int(image_id)
|
||||
except (TypeError, ValueError):
|
||||
return None, None, self._bad_request("image_id must be an integer")
|
||||
record = await self._run_in_executor(self.db.get_image_by_id, image_id)
|
||||
if not record or not record.get("image_path"):
|
||||
return (
|
||||
None,
|
||||
None,
|
||||
web.json_response(
|
||||
{"success": False, "error": f"Image {image_id} not found"},
|
||||
status=404,
|
||||
),
|
||||
)
|
||||
stored_path = record["image_path"]
|
||||
resolved = self._resolve_allowed_image_path(stored_path)
|
||||
if resolved is None:
|
||||
return None, None, self._forbidden(stored_path)
|
||||
return resolved, record.get("prompt_id"), None
|
||||
|
||||
raw_path = data.get("path") or data.get("image_path")
|
||||
if not raw_path or not isinstance(raw_path, str):
|
||||
return None, None, self._bad_request("image_id or path is required")
|
||||
|
||||
if "root_index" in data:
|
||||
root_index = data["root_index"]
|
||||
if isinstance(root_index, bool) or not isinstance(root_index, int):
|
||||
return None, None, self._bad_request("root_index must be an integer")
|
||||
roots = self._get_all_output_dirs()
|
||||
if not 0 <= root_index < len(roots):
|
||||
return None, None, self._forbidden(raw_path)
|
||||
root = roots[root_index]
|
||||
if Path(raw_path).is_absolute() or ".." in raw_path.replace(
|
||||
"\\", "/"
|
||||
).split("/"):
|
||||
return None, None, self._forbidden(raw_path)
|
||||
raw_path = str(Path(root) / raw_path)
|
||||
if not path_is_within(raw_path, root):
|
||||
return None, None, self._forbidden(raw_path)
|
||||
|
||||
resolved = self._resolve_allowed_image_path(raw_path)
|
||||
if resolved is None:
|
||||
return None, None, self._forbidden(raw_path)
|
||||
|
||||
prompt_id = None
|
||||
try:
|
||||
prompt_id = await self._run_in_executor(
|
||||
self.db.get_prompt_id_for_image, raw_path
|
||||
)
|
||||
if prompt_id is None and resolved != raw_path:
|
||||
prompt_id = await self._run_in_executor(
|
||||
self.db.get_prompt_id_for_image, resolved
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not find linked prompt: {e}")
|
||||
return resolved, prompt_id, None
|
||||
|
||||
@staticmethod
|
||||
def _bad_request(message: str):
|
||||
return web.json_response({"success": False, "error": message}, status=400)
|
||||
|
||||
async def autotag_single(self, request):
|
||||
"""Generate tags for a single image selected by image_id or gallery path."""
|
||||
try:
|
||||
data = await request.json()
|
||||
if not isinstance(data, dict):
|
||||
return self._bad_request("Request body must be a JSON object")
|
||||
model_type = data.get("model_type", "gguf")
|
||||
custom_prompt = data.get("prompt")
|
||||
use_gpu = data.get("use_gpu", True)
|
||||
general_threshold = data.get("general_threshold")
|
||||
character_threshold = data.get("character_threshold")
|
||||
try:
|
||||
if general_threshold is not None:
|
||||
general_threshold = float(general_threshold)
|
||||
if character_threshold is not None:
|
||||
character_threshold = float(character_threshold)
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
)
|
||||
|
||||
image_path, prompt_id, error = await self._select_autotag_target(data)
|
||||
if error is not None:
|
||||
return error
|
||||
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
service = get_autotag_service()
|
||||
|
||||
if (
|
||||
not service.is_model_loaded()
|
||||
or service.get_loaded_model_type() != model_type
|
||||
):
|
||||
await self._run_in_executor(service.load_model, model_type, use_gpu)
|
||||
|
||||
if custom_prompt:
|
||||
service.custom_prompt = custom_prompt
|
||||
|
||||
tags = await self._run_in_executor(
|
||||
service.generate_tags,
|
||||
image_path,
|
||||
general_threshold=general_threshold,
|
||||
character_threshold=character_threshold,
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"tags": tags,
|
||||
"prompt_id": prompt_id,
|
||||
"filename": os.path.basename(image_path),
|
||||
}
|
||||
)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
return self._bad_request("Request body must be valid JSON")
|
||||
except Exception as e:
|
||||
self.logger.error(f"AutoTag single error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
|
||||
async def apply_autotag(self, request):
|
||||
"""Apply selected tags to a prompt."""
|
||||
try:
|
||||
data = await request.json()
|
||||
prompt_id = data.get("prompt_id")
|
||||
tags = data.get("tags", [])
|
||||
|
||||
if not prompt_id:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "prompt_id is required"}, status=400
|
||||
)
|
||||
|
||||
if not tags:
|
||||
return web.json_response(
|
||||
{"success": True, "message": "No tags to apply"}
|
||||
)
|
||||
|
||||
prompt = await self._run_in_executor(self.db.get_prompt_by_id, prompt_id)
|
||||
if not prompt:
|
||||
return web.json_response(
|
||||
{"success": False, "error": f"Prompt {prompt_id} not found"},
|
||||
status=404,
|
||||
)
|
||||
|
||||
existing_tags = prompt.get("tags", [])
|
||||
if isinstance(existing_tags, str):
|
||||
existing_tags = [
|
||||
t.strip() for t in existing_tags.split(",") if t.strip()
|
||||
]
|
||||
|
||||
new_tags = [t for t in tags if t not in existing_tags]
|
||||
all_tags = existing_tags + new_tags
|
||||
|
||||
await self._run_in_executor(
|
||||
self.db.update_prompt_metadata, prompt_id, tags=all_tags
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{"success": True, "added_tags": new_tags, "total_tags": len(all_tags)}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Apply autotag error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
|
||||
async def unload_autotag_model(self, request):
|
||||
"""Manually unload the AutoTag model from memory."""
|
||||
try:
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
service = get_autotag_service()
|
||||
|
||||
if not service.is_model_loaded():
|
||||
return web.json_response(
|
||||
{"success": True, "message": "No model was loaded"}
|
||||
)
|
||||
|
||||
model_type = service.get_loaded_model_type()
|
||||
service.unload_model()
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"message": f"{model_type.upper()} model unloaded successfully",
|
||||
"model_loaded": False,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Unload autotag model error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _collect_scan_files(output_path: Path):
|
||||
"""Media files under ``output_path`` from a bounded, symlink-free walk.
|
||||
|
||||
Descends at most ``_SCAN_MAX_DEPTH`` levels, skips ``thumbnails``
|
||||
trees and symlinked files, and stops after ``_SCAN_MAX_FILES`` files.
|
||||
Returns ``(paths, truncated)``.
|
||||
"""
|
||||
found = []
|
||||
base_depth = len(output_path.parts)
|
||||
for dirpath, dirnames, filenames in os.walk(output_path):
|
||||
depth = len(Path(dirpath).parts) - base_depth
|
||||
dirnames[:] = sorted(
|
||||
name
|
||||
for name in dirnames
|
||||
if name != "thumbnails" and depth + 1 <= _SCAN_MAX_DEPTH
|
||||
)
|
||||
for name in sorted(filenames):
|
||||
if Path(name).suffix.lower() not in _MEDIA_EXTENSIONS:
|
||||
continue
|
||||
candidate = Path(dirpath) / name
|
||||
if candidate.is_symlink():
|
||||
continue
|
||||
if len(found) >= _SCAN_MAX_FILES:
|
||||
return found, True
|
||||
found.append(candidate)
|
||||
return found, False
|
||||
|
||||
async def scan_output_dir(self, request):
|
||||
"""Scan ComfyUI output directory for images."""
|
||||
try:
|
||||
output_dir = self._find_comfyui_output_dir()
|
||||
if not output_dir:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "ComfyUI output directory not found"},
|
||||
status=404,
|
||||
)
|
||||
|
||||
output_path = Path(output_dir)
|
||||
files, truncated = await self._run_in_executor(
|
||||
self._collect_scan_files, output_path
|
||||
)
|
||||
has_thumbnails = (output_path / "thumbnails").is_dir()
|
||||
|
||||
images = []
|
||||
for image_path in files:
|
||||
rel_path = image_path.relative_to(output_path)
|
||||
thumbnail_url = None
|
||||
if has_thumbnails:
|
||||
thumb_rel = (
|
||||
f"thumbnails/{rel_path.with_suffix('').as_posix()}"
|
||||
f"_thumb{image_path.suffix}"
|
||||
)
|
||||
thumbnail_url = thumbnail_url_for(
|
||||
thumb_rel, output_path / thumb_rel
|
||||
)
|
||||
images.append(
|
||||
{
|
||||
"filename": image_path.name,
|
||||
"path": self._public_path(image_path),
|
||||
"relative_path": rel_path.as_posix(),
|
||||
"url": "/prompt_manager/images/serve/"
|
||||
+ quote(rel_path.as_posix(), safe="/"),
|
||||
"thumbnail_url": thumbnail_url,
|
||||
}
|
||||
)
|
||||
|
||||
images.sort(key=lambda x: x["filename"])
|
||||
|
||||
self.logger.info(f"Found {len(images)} images in output directory")
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"images": images,
|
||||
"count": len(images),
|
||||
"truncated": truncated,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Scan output dir error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
@@ -0,0 +1,299 @@
|
||||
"""Logging API routes for PromptManager."""
|
||||
|
||||
import os
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
# Keys PromptManagerLogger.update_config() understands. Anything else is
|
||||
# rejected so a client cannot inject arbitrary entries into the logging config.
|
||||
LOG_CONFIG_KEYS = frozenset(
|
||||
{
|
||||
"level",
|
||||
"max_file_size",
|
||||
"backup_count",
|
||||
"console_logging",
|
||||
"file_logging",
|
||||
"buffer_size",
|
||||
}
|
||||
)
|
||||
_LOG_INT_KEYS = frozenset({"max_file_size", "backup_count", "buffer_size"})
|
||||
_LOG_BOOL_KEYS = frozenset({"console_logging", "file_logging"})
|
||||
_LOG_LEVEL_NAMES = frozenset({"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"})
|
||||
# Inclusive bounds for the integer keys that size the log files on disk
|
||||
_MEGABYTE = 1024 * 1024
|
||||
_LOG_INT_BOUNDS = {
|
||||
"max_file_size": (1 * _MEGABYTE, 1024 * _MEGABYTE),
|
||||
"backup_count": (0, 50),
|
||||
}
|
||||
_LOG_INT_BOUND_LABELS = {
|
||||
"max_file_size": "between 1 MB and 1024 MB (in bytes)",
|
||||
"backup_count": "between 0 and 50",
|
||||
}
|
||||
|
||||
DEFAULT_LOG_LIMIT = 100
|
||||
MAX_LOG_LIMIT = 500
|
||||
|
||||
|
||||
def _parse_limit(raw):
|
||||
"""Parse a ``limit`` query value, clamped to [1, MAX_LOG_LIMIT].
|
||||
|
||||
Raises:
|
||||
ValueError: when ``raw`` is not an integer.
|
||||
"""
|
||||
if raw is None or raw == "":
|
||||
return DEFAULT_LOG_LIMIT
|
||||
limit = int(raw)
|
||||
return max(1, min(limit, MAX_LOG_LIMIT))
|
||||
|
||||
|
||||
def _validate_log_config(data):
|
||||
"""Validate a log-config payload against LOG_CONFIG_KEYS.
|
||||
|
||||
Returns:
|
||||
(clean, error): ``clean`` is the normalised dict to apply, ``error``
|
||||
a message describing the first problem (``clean`` is then ``None``).
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
return None, "Request body must be a JSON object"
|
||||
|
||||
unknown = sorted(key for key in data if key not in LOG_CONFIG_KEYS)
|
||||
if unknown:
|
||||
return None, (
|
||||
f"Unknown logging config keys: {', '.join(unknown)}. "
|
||||
f"Allowed: {', '.join(sorted(LOG_CONFIG_KEYS))}"
|
||||
)
|
||||
|
||||
clean = {}
|
||||
for key, value in data.items():
|
||||
if key == "level":
|
||||
if not isinstance(value, str) or value.upper() not in _LOG_LEVEL_NAMES:
|
||||
return None, (
|
||||
f"Invalid log level. Must be one of: {sorted(_LOG_LEVEL_NAMES)}"
|
||||
)
|
||||
clean[key] = value.upper()
|
||||
elif key in _LOG_INT_KEYS:
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
||||
return None, f"{key} must be a non-negative integer"
|
||||
bounds = _LOG_INT_BOUNDS.get(key)
|
||||
if bounds and not bounds[0] <= value <= bounds[1]:
|
||||
return None, f"{key} must be {_LOG_INT_BOUND_LABELS[key]}"
|
||||
clean[key] = value
|
||||
elif key in _LOG_BOOL_KEYS:
|
||||
if not isinstance(value, bool):
|
||||
return None, f"{key} must be a boolean"
|
||||
clean[key] = value
|
||||
return clean, None
|
||||
|
||||
|
||||
class LoggingRoutesMixin:
|
||||
"""Mixin providing logging-related API endpoints."""
|
||||
|
||||
def _register_logging_routes(self, routes):
|
||||
@routes.get("/prompt_manager/logs")
|
||||
async def get_logs_route(request):
|
||||
return await self.get_logs(request)
|
||||
|
||||
@routes.get("/prompt_manager/logs/files")
|
||||
async def get_log_files_route(request):
|
||||
return await self.get_log_files(request)
|
||||
|
||||
@routes.get("/prompt_manager/logs/download/{filename}")
|
||||
async def download_log_route(request):
|
||||
return await self.download_log_file(request)
|
||||
|
||||
@routes.post("/prompt_manager/logs/truncate")
|
||||
async def truncate_logs_route(request):
|
||||
return await self.truncate_logs(request)
|
||||
|
||||
@routes.get("/prompt_manager/logs/config")
|
||||
async def get_log_config_route(request):
|
||||
return await self.get_log_config(request)
|
||||
|
||||
@routes.post("/prompt_manager/logs/config")
|
||||
async def update_log_config_route(request):
|
||||
return await self.update_log_config(request)
|
||||
|
||||
@routes.get("/prompt_manager/logs/stats")
|
||||
async def get_log_stats_route(request):
|
||||
return await self.get_log_stats(request)
|
||||
|
||||
def _get_logger_manager(self):
|
||||
"""Get the logger manager instance."""
|
||||
try:
|
||||
from ...utils.logging_config import get_logger_manager
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
current_dir = os.path.dirname(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
)
|
||||
sys.path.insert(0, current_dir)
|
||||
from utils.logging_config import get_logger_manager
|
||||
return get_logger_manager()
|
||||
|
||||
async def get_logs(self, request):
|
||||
"""Get recent log entries."""
|
||||
try:
|
||||
logger_manager = self._get_logger_manager()
|
||||
|
||||
try:
|
||||
limit = _parse_limit(request.query.get("limit"))
|
||||
except ValueError:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "limit must be an integer", "logs": []},
|
||||
status=400,
|
||||
)
|
||||
level = request.query.get("level", None)
|
||||
|
||||
logs = logger_manager.get_recent_logs(limit=limit, level=level)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"logs": logs,
|
||||
"count": len(logs),
|
||||
"level_filter": level,
|
||||
"limit": limit,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Get logs error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e), "logs": []},
|
||||
status=500,
|
||||
)
|
||||
|
||||
async def get_log_files(self, request):
|
||||
"""Get information about available log files."""
|
||||
try:
|
||||
logger_manager = self._get_logger_manager()
|
||||
log_files = [
|
||||
{**entry, "path": self._public_path(entry.get("path"))}
|
||||
for entry in logger_manager.get_log_files()
|
||||
]
|
||||
|
||||
return web.json_response(
|
||||
{"success": True, "files": log_files, "count": len(log_files)}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Get log files error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e), "files": []},
|
||||
status=500,
|
||||
)
|
||||
|
||||
async def download_log_file(self, request):
|
||||
"""Download a specific log file."""
|
||||
try:
|
||||
filename = request.match_info["filename"]
|
||||
logger_manager = self._get_logger_manager()
|
||||
|
||||
if not filename or ".." in filename or "/" in filename or "\\" in filename:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Invalid filename"}, status=400
|
||||
)
|
||||
|
||||
log_file_path = logger_manager.log_dir / filename
|
||||
|
||||
if not log_file_path.exists():
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Log file not found"}, status=404
|
||||
)
|
||||
|
||||
with open(log_file_path, "rb") as f:
|
||||
file_content = f.read()
|
||||
|
||||
return web.Response(
|
||||
body=file_content,
|
||||
content_type="text/plain",
|
||||
headers={
|
||||
"Content-Disposition": f'attachment; filename="{filename}"',
|
||||
"Content-Length": str(len(file_content)),
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Download log file error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
|
||||
async def truncate_logs(self, request):
|
||||
"""Truncate all log files."""
|
||||
try:
|
||||
logger_manager = self._get_logger_manager()
|
||||
results = logger_manager.truncate_logs()
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"message": f"Truncated {len(results['truncated'])} log files",
|
||||
"results": results,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Truncate logs error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
|
||||
async def get_log_config(self, request):
|
||||
"""Get current logging configuration."""
|
||||
try:
|
||||
logger_manager = self._get_logger_manager()
|
||||
config = logger_manager.get_config()
|
||||
|
||||
return web.json_response({"success": True, "config": config})
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Get log config error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
|
||||
async def update_log_config(self, request):
|
||||
"""Update logging configuration (whitelisted keys only)."""
|
||||
try:
|
||||
data, error_response = await self._read_json_body(request)
|
||||
if error_response is not None:
|
||||
return error_response
|
||||
|
||||
clean, error = _validate_log_config(data)
|
||||
if error:
|
||||
return web.json_response({"success": False, "error": error}, status=400)
|
||||
|
||||
logger_manager = self._get_logger_manager()
|
||||
logger_manager.update_config(clean)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"message": "Logging configuration updated",
|
||||
"config": logger_manager.get_config(),
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Update log config error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Failed to update logging configuration"},
|
||||
status=500,
|
||||
)
|
||||
|
||||
async def get_log_stats(self, request):
|
||||
"""Get logging statistics."""
|
||||
try:
|
||||
logger_manager = self._get_logger_manager()
|
||||
stats = dict(logger_manager.get_log_stats())
|
||||
stats["log_directory"] = self._public_path(stats.get("log_directory"))
|
||||
|
||||
return web.json_response({"success": True, "stats": stats})
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Get log stats error: {e}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
@@ -0,0 +1,456 @@
|
||||
"""LoraManager integration API routes for PromptManager."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
_LORA_PATH_ERROR = (
|
||||
"LoraManager path must be a directory under ComfyUI's custom_nodes "
|
||||
"whose name contains 'lora'"
|
||||
)
|
||||
|
||||
|
||||
class LoraIntegrationMixin:
|
||||
"""Mixin providing LoraManager detection, scanning, and trigger word endpoints."""
|
||||
|
||||
def _register_lora_routes(self, routes):
|
||||
@routes.get("/prompt_manager/lora/detect")
|
||||
async def lora_detect_route(request):
|
||||
return await self.lora_detect(request)
|
||||
|
||||
@routes.get("/prompt_manager/lora/status")
|
||||
async def lora_status_route(request):
|
||||
return await self.lora_status(request)
|
||||
|
||||
@routes.post("/prompt_manager/lora/enable")
|
||||
async def lora_enable_route(request):
|
||||
return await self.lora_enable(request)
|
||||
|
||||
@routes.post("/prompt_manager/lora/scan")
|
||||
async def lora_scan_route(request):
|
||||
return await self.lora_scan(request)
|
||||
|
||||
@routes.get("/prompt_manager/lora/trigger-words")
|
||||
async def lora_trigger_words_route(request):
|
||||
return await self.lora_trigger_words(request)
|
||||
|
||||
@routes.post("/prompt_manager/lora/refresh-cache")
|
||||
async def lora_refresh_cache_route(request):
|
||||
return await self.lora_refresh_cache(request)
|
||||
|
||||
def _error_response(self, exc: Exception, status: int = 500):
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(exc)}, status=status
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _bad_lora_request(message: str):
|
||||
return web.json_response({"success": False, "error": message}, status=400)
|
||||
|
||||
# ── Detection ────────────────────────────────────────────────────
|
||||
|
||||
async def lora_detect(self, request):
|
||||
"""Auto-detect LoraManager installation."""
|
||||
try:
|
||||
from ..lora_utils import detect_lora_manager
|
||||
|
||||
path = await self._run_in_executor(detect_lora_manager)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"detected": path is not None,
|
||||
"path": self._public_path(path) if path else "",
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"LoraManager detection failed: {e}")
|
||||
return self._error_response(e)
|
||||
|
||||
# ── Status ───────────────────────────────────────────────────────
|
||||
|
||||
async def lora_status(self, request):
|
||||
"""Get current LoraManager integration status."""
|
||||
try:
|
||||
from ..config import IntegrationConfig
|
||||
from ..lora_utils import detect_lora_manager, get_trigger_cache
|
||||
|
||||
config = IntegrationConfig.get_config()["lora_manager"]
|
||||
cache = get_trigger_cache()
|
||||
|
||||
# Check if the configured path is still valid
|
||||
detected_path = await self._run_in_executor(
|
||||
detect_lora_manager, config.get("path", "")
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"enabled": config["enabled"],
|
||||
"path": self._public_path(config["path"]),
|
||||
"trigger_words_enabled": config["trigger_words_enabled"],
|
||||
"has_civitai_api_key": bool(config.get("civitai_api_key")),
|
||||
"detected": detected_path is not None,
|
||||
"detected_path": (
|
||||
self._public_path(detected_path) if detected_path else ""
|
||||
),
|
||||
"trigger_cache_loaded": cache.is_loaded,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"LoraManager status check failed: {e}")
|
||||
return self._error_response(e)
|
||||
|
||||
# ── Enable / Disable ─────────────────────────────────────────────
|
||||
|
||||
async def _validated_lora_manager_path(self, path):
|
||||
"""Canonical LoraManager path for ``path`` or ``(None, error_response)``.
|
||||
|
||||
The user's path (or, when empty, the auto-detected install) must be a
|
||||
LoraManager-looking directory that is a direct child of ComfyUI's
|
||||
``custom_nodes`` with ``lora`` in its name; the trees under it are
|
||||
later served as images, so nothing outside custom_nodes may be used.
|
||||
"""
|
||||
from ..lora_utils import detect_lora_manager, resolve_lora_manager_path
|
||||
|
||||
candidate = ""
|
||||
if path:
|
||||
candidate = await self._run_in_executor(resolve_lora_manager_path, path)
|
||||
if candidate is None:
|
||||
return None, self._bad_lora_request(_LORA_PATH_ERROR)
|
||||
|
||||
resolved = await self._run_in_executor(detect_lora_manager, candidate)
|
||||
if not resolved:
|
||||
return None, self._bad_lora_request(
|
||||
"LoraManager not found at the specified path"
|
||||
)
|
||||
|
||||
allowed = await self._run_in_executor(resolve_lora_manager_path, resolved)
|
||||
if allowed is None:
|
||||
return None, self._bad_lora_request(_LORA_PATH_ERROR)
|
||||
return allowed, None
|
||||
|
||||
async def lora_enable(self, request):
|
||||
"""Enable or disable LoraManager integration and save to config.json."""
|
||||
try:
|
||||
data = await request.json()
|
||||
enabled = data.get("enabled", False)
|
||||
path = data.get("path", "")
|
||||
trigger_words = data.get("trigger_words_enabled", False)
|
||||
|
||||
from ..config import IntegrationConfig, PromptManagerConfig
|
||||
from ..lora_utils import get_trigger_cache
|
||||
|
||||
# The key is never echoed back to the UI, so a blank or absent
|
||||
# field means "unchanged"; clearing is an explicit request.
|
||||
civitai_key = IntegrationConfig.CIVITAI_API_KEY
|
||||
submitted_key = data.get("civitai_api_key")
|
||||
if submitted_key is not None and not isinstance(submitted_key, str):
|
||||
return self._bad_lora_request("civitai_api_key must be a string")
|
||||
if data.get("clear_civitai_api_key"):
|
||||
civitai_key = ""
|
||||
elif submitted_key and submitted_key.strip():
|
||||
civitai_key = submitted_key.strip()
|
||||
|
||||
if not isinstance(path, str):
|
||||
return self._bad_lora_request("path must be a string")
|
||||
|
||||
# If enabling, validate the path
|
||||
if enabled:
|
||||
path, error = await self._validated_lora_manager_path(path)
|
||||
if error is not None:
|
||||
return error
|
||||
|
||||
# Update in-memory config
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = enabled
|
||||
IntegrationConfig.LORA_MANAGER_PATH = path
|
||||
IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = trigger_words
|
||||
IntegrationConfig.CIVITAI_API_KEY = civitai_key
|
||||
|
||||
# Persist to the configured config.json
|
||||
PromptManagerConfig.save_to_file()
|
||||
|
||||
# Load trigger word cache if enabling
|
||||
cache = get_trigger_cache()
|
||||
if enabled and trigger_words and path:
|
||||
count = await self._run_in_executor(cache.load, path)
|
||||
self.logger.info(f"Trigger word cache loaded: {count} LoRAs")
|
||||
elif not enabled:
|
||||
cache.clear()
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"enabled": enabled,
|
||||
"path": self._public_path(path),
|
||||
"trigger_words_enabled": trigger_words,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"LoraManager enable/disable failed: {e}")
|
||||
return self._error_response(e)
|
||||
|
||||
# ── Scan LoRA example images ─────────────────────────────────────
|
||||
|
||||
async def lora_scan(self, request):
|
||||
"""Scan LoraManager metadata and import LoRA info + preview images.
|
||||
|
||||
Streams progress as SSE, matching the existing scan pattern. Errors
|
||||
raised once the stream is open are reported as an ``error`` event.
|
||||
"""
|
||||
try:
|
||||
from ..config import IntegrationConfig
|
||||
from ..lora_utils import resolve_lora_manager_path
|
||||
|
||||
if not IntegrationConfig.LORA_MANAGER_ENABLED:
|
||||
return self._bad_lora_request("LoraManager integration is not enabled")
|
||||
|
||||
lm_path = IntegrationConfig.LORA_MANAGER_PATH
|
||||
if not lm_path:
|
||||
return self._bad_lora_request("LoraManager path not configured")
|
||||
|
||||
allowed = await self._run_in_executor(resolve_lora_manager_path, lm_path)
|
||||
if allowed is None:
|
||||
return self._bad_lora_request(_LORA_PATH_ERROR)
|
||||
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
reason="OK",
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
await response.prepare(request)
|
||||
|
||||
async def send_progress(data):
|
||||
line = f"data: {json.dumps(data)}\n\n"
|
||||
await response.write(line.encode("utf-8"))
|
||||
|
||||
try:
|
||||
await self._lora_scan_stream(allowed, send_progress)
|
||||
except Exception as e:
|
||||
self.logger.exception("LoRA scan failed")
|
||||
await send_progress({"type": "error", "message": self._public_error(e)})
|
||||
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"LoRA scan failed: {e}")
|
||||
return self._error_response(e)
|
||||
|
||||
async def _lora_scan_stream(self, lm_path, send_progress):
|
||||
"""Import every LoRA under ``lm_path``, reporting through ``send_progress``."""
|
||||
from ..config import IntegrationConfig
|
||||
from ..lora_utils import (
|
||||
download_civitai_images,
|
||||
find_lora_directories,
|
||||
get_example_prompt_from_metadata,
|
||||
get_lora_image_cache_dir,
|
||||
get_preview_images_from_metadata,
|
||||
get_trigger_words_from_metadata,
|
||||
get_model_name_from_metadata,
|
||||
read_lora_metadata,
|
||||
)
|
||||
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
"status": "Clearing previous lora-manager imports...",
|
||||
"progress": 0,
|
||||
}
|
||||
)
|
||||
|
||||
# Clear previous imports so reimport is always clean
|
||||
await self._run_in_executor(self.db.delete_prompts_by_category, "lora-manager")
|
||||
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
"status": "Finding LoRA directories...",
|
||||
"progress": 2,
|
||||
}
|
||||
)
|
||||
|
||||
lora_dirs = await self._run_in_executor(find_lora_directories, lm_path)
|
||||
|
||||
# Collect all metadata files
|
||||
meta_files = []
|
||||
for d in lora_dirs:
|
||||
dir_path = Path(d)
|
||||
meta_files.extend(dir_path.rglob("*.metadata.json"))
|
||||
|
||||
total = len(meta_files)
|
||||
imported = 0
|
||||
skipped = 0
|
||||
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
"status": f"Found {total} LoRA metadata files",
|
||||
"progress": 5,
|
||||
"total": total,
|
||||
}
|
||||
)
|
||||
|
||||
cache_dir = get_lora_image_cache_dir()
|
||||
|
||||
for i, meta_file in enumerate(meta_files):
|
||||
metadata = await self._run_in_executor(read_lora_metadata, meta_file)
|
||||
if not metadata:
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
model_name = get_model_name_from_metadata(metadata)
|
||||
trigger_words = get_trigger_words_from_metadata(metadata)
|
||||
|
||||
# Collect all images: local previews + downloaded civitai examples
|
||||
preview_paths = await self._run_in_executor(
|
||||
get_preview_images_from_metadata, metadata, meta_file
|
||||
)
|
||||
civitai_paths = await self._run_in_executor(
|
||||
download_civitai_images,
|
||||
metadata,
|
||||
meta_file,
|
||||
cache_dir,
|
||||
IntegrationConfig.CIVITAI_API_KEY,
|
||||
)
|
||||
|
||||
# Merge, local first, dedup
|
||||
seen = set(preview_paths)
|
||||
all_images = list(preview_paths)
|
||||
for cp in civitai_paths:
|
||||
if cp not in seen:
|
||||
all_images.append(cp)
|
||||
seen.add(cp)
|
||||
|
||||
# Build prompt text: prefer example prompt, then model name
|
||||
example_prompt = get_example_prompt_from_metadata(metadata)
|
||||
prompt_text = example_prompt or model_name
|
||||
|
||||
# Build tags
|
||||
tags = ["lora-manager", f"lora:{model_name}"]
|
||||
tags.extend(trigger_words)
|
||||
|
||||
# Save to database via existing mechanism
|
||||
try:
|
||||
import hashlib
|
||||
|
||||
prompt_hash = hashlib.sha256(
|
||||
prompt_text.strip().lower().encode("utf-8")
|
||||
).hexdigest()
|
||||
|
||||
existing = await self._run_in_executor(
|
||||
self.db.get_prompt_by_hash, prompt_hash
|
||||
)
|
||||
|
||||
if existing:
|
||||
# Link all images
|
||||
for pp in all_images:
|
||||
await self._run_in_executor(
|
||||
self.db.link_image_to_prompt,
|
||||
existing["id"],
|
||||
pp,
|
||||
)
|
||||
skipped += 1
|
||||
else:
|
||||
prompt_id = await self._run_in_executor(
|
||||
self.db.save_prompt,
|
||||
prompt_text,
|
||||
"lora-manager", # category
|
||||
tags,
|
||||
None, # rating
|
||||
None, # notes
|
||||
prompt_hash,
|
||||
)
|
||||
|
||||
if prompt_id:
|
||||
for pp in all_images:
|
||||
await self._run_in_executor(
|
||||
self.db.link_image_to_prompt,
|
||||
prompt_id,
|
||||
pp,
|
||||
)
|
||||
imported += 1
|
||||
else:
|
||||
skipped += 1
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to import LoRA {model_name}: {e}")
|
||||
skipped += 1
|
||||
|
||||
# Progress update for every LoRA
|
||||
progress = int(5 + (90 * (i + 1) / max(total, 1)))
|
||||
img_count = len(all_images)
|
||||
status = f"{model_name}"
|
||||
if img_count:
|
||||
status += f" ({img_count} images)"
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
"status": status,
|
||||
"progress": progress,
|
||||
"processed": i + 1,
|
||||
"imported": imported,
|
||||
"skipped": skipped,
|
||||
}
|
||||
)
|
||||
|
||||
await send_progress(
|
||||
{
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
"total": total,
|
||||
"imported": imported,
|
||||
"skipped": skipped,
|
||||
}
|
||||
)
|
||||
|
||||
# ── Trigger word endpoints ───────────────────────────────────────
|
||||
|
||||
async def lora_trigger_words(self, request):
|
||||
"""Look up trigger words for a specific LoRA name."""
|
||||
try:
|
||||
from ..config import IntegrationConfig
|
||||
from ..lora_utils import get_trigger_cache
|
||||
|
||||
if not IntegrationConfig.LORA_MANAGER_ENABLED:
|
||||
return self._bad_lora_request("LoraManager integration not enabled")
|
||||
|
||||
lora_name = request.query.get("name", "")
|
||||
if not lora_name:
|
||||
return self._bad_lora_request("Missing 'name' query parameter")
|
||||
|
||||
cache = get_trigger_cache()
|
||||
if not cache.is_loaded:
|
||||
lm_path = IntegrationConfig.LORA_MANAGER_PATH
|
||||
if lm_path:
|
||||
await self._run_in_executor(cache.load, lm_path)
|
||||
|
||||
words = cache.get_trigger_words(lora_name)
|
||||
return web.json_response(
|
||||
{"success": True, "lora": lora_name, "trigger_words": words}
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Trigger word lookup failed: {e}")
|
||||
return self._error_response(e)
|
||||
|
||||
async def lora_refresh_cache(self, request):
|
||||
"""Force-refresh the trigger word cache from disk."""
|
||||
try:
|
||||
from ..config import IntegrationConfig
|
||||
from ..lora_utils import get_trigger_cache
|
||||
|
||||
if not IntegrationConfig.LORA_MANAGER_ENABLED:
|
||||
return self._bad_lora_request("LoraManager integration not enabled")
|
||||
|
||||
lm_path = IntegrationConfig.LORA_MANAGER_PATH
|
||||
if not lm_path:
|
||||
return self._bad_lora_request("LoraManager path not configured")
|
||||
|
||||
cache = get_trigger_cache()
|
||||
count = await self._run_in_executor(cache.load, lm_path)
|
||||
return web.json_response(
|
||||
{"success": True, "loras_with_trigger_words": count}
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Trigger cache refresh failed: {e}")
|
||||
return self._error_response(e)
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Progress fan-out for the background output scan.
|
||||
|
||||
The scan runs as an asyncio task that outlives the HTTP request that started
|
||||
it, so an accidental tab close does not abandon a half-finished scan. Every
|
||||
SSE subscriber (the request that started the scan, a reconnecting page, a
|
||||
second tab) receives the latest event as a snapshot and then every event
|
||||
published after it, until the job finishes.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
|
||||
class ScanJob:
|
||||
"""One scan run: its latest event, completion flag and live subscribers."""
|
||||
|
||||
def __init__(self):
|
||||
self.last_event = None
|
||||
self.finished = False
|
||||
self.task = None
|
||||
self._subscribers = set()
|
||||
|
||||
@property
|
||||
def running(self):
|
||||
return not self.finished
|
||||
|
||||
@property
|
||||
def subscriber_count(self):
|
||||
return len(self._subscribers)
|
||||
|
||||
def publish(self, payload):
|
||||
"""Record ``payload`` as the latest event and hand it to every subscriber."""
|
||||
self.last_event = payload
|
||||
for queue in list(self._subscribers):
|
||||
queue.put_nowait(payload)
|
||||
|
||||
def finish(self):
|
||||
"""Mark the job done and release every subscriber after its queued events."""
|
||||
self.finished = True
|
||||
for queue in list(self._subscribers):
|
||||
queue.put_nowait(None)
|
||||
|
||||
async def events(self):
|
||||
"""Yield the latest event, then each later one, until the job finishes.
|
||||
|
||||
Consume it under ``contextlib.aclosing`` so a consumer that stops early
|
||||
(a closed socket) drops its queue at once instead of at garbage collection.
|
||||
"""
|
||||
queue = asyncio.Queue()
|
||||
self._subscribers.add(queue)
|
||||
try:
|
||||
if self.last_event is not None:
|
||||
yield self.last_event
|
||||
if self.finished:
|
||||
return
|
||||
while True:
|
||||
payload = await queue.get()
|
||||
if payload is None:
|
||||
return
|
||||
yield payload
|
||||
finally:
|
||||
self._subscribers.discard(queue)
|
||||
@@ -0,0 +1,945 @@
|
||||
"""
|
||||
AutoTag Service Module
|
||||
|
||||
Provides JoyCaption-based automatic tagging for images using LLM models.
|
||||
Refactored from standalone_tagger.py for integration with PromptManager API.
|
||||
"""
|
||||
|
||||
import errno
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
from PIL import Image
|
||||
|
||||
# A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend)
|
||||
ModelLoader = Callable[[str, bool], Tuple[str, Any]]
|
||||
|
||||
# Try to import logging from utils, fallback to standard logging
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import logging
|
||||
|
||||
def get_logger(name: str) -> logging.Logger:
|
||||
logger = logging.getLogger(name)
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler()
|
||||
handler.setFormatter(
|
||||
logging.Formatter(
|
||||
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
)
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
return logger
|
||||
|
||||
|
||||
# Model configurations - same as standalone_tagger.py
|
||||
MODELS = {
|
||||
"gguf": {
|
||||
"name": "GGUF (Recommended)",
|
||||
"description": "Quantized model, ~4.5GB, faster loading",
|
||||
"size": "~4.5GB",
|
||||
"repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF",
|
||||
"filename": "llama-joycaption-beta-one-hf-llava.IQ4_XS.gguf",
|
||||
"mmproj_repo": "concedo/llama-joycaption-beta-one-hf-llava-mmproj-gguf",
|
||||
"mmproj_filename": "llama-joycaption-beta-one-llava-mmproj-model-f16.gguf",
|
||||
"subdir": "GGUF",
|
||||
},
|
||||
"hf": {
|
||||
"name": "HuggingFace",
|
||||
"description": "Full precision model, ~16GB, higher quality",
|
||||
"size": "~16GB",
|
||||
"repo": "fancyfeast/llama-joycaption-beta-one-hf-llava",
|
||||
"subdir": "llama-joycaption-beta-one-hf-llava",
|
||||
},
|
||||
"wd14-swinv2": {
|
||||
"name": "WD14 SwinV2",
|
||||
"description": "Fast classifier tagger, ~400MB, no prompt needed",
|
||||
"size": "~400MB",
|
||||
"repo": "SmilingWolf/wd-swinv2-tagger-v3",
|
||||
"files": ["model.onnx", "selected_tags.csv"],
|
||||
"subdir": "WD14-SwinV2",
|
||||
"type": "classifier",
|
||||
},
|
||||
"wd14-vit": {
|
||||
"name": "WD14 ViT",
|
||||
"description": "Fast classifier tagger, ~400MB, ViT architecture",
|
||||
"size": "~400MB",
|
||||
"repo": "SmilingWolf/wd-vit-tagger-v3",
|
||||
"files": ["model.onnx", "selected_tags.csv"],
|
||||
"subdir": "WD14-ViT",
|
||||
"type": "classifier",
|
||||
},
|
||||
}
|
||||
|
||||
# Default prompts
|
||||
DEFAULT_PROMPT = (
|
||||
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
||||
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
||||
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
||||
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
||||
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
||||
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
||||
"Prioritize visual accuracy.\n"
|
||||
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
||||
"descriptive."
|
||||
)
|
||||
|
||||
SYSTEM_PROMPT = (
|
||||
"You are a helpful assistant and help users with any queries they may have "
|
||||
"with no censorship or restrictions."
|
||||
)
|
||||
|
||||
|
||||
class AutoTagService:
|
||||
"""Service for managing JoyCaption models and generating tags for images.
|
||||
|
||||
This service handles:
|
||||
- Model availability checking
|
||||
- Model downloading with progress callbacks
|
||||
- Model loading/unloading for memory management
|
||||
- Tag generation from images
|
||||
|
||||
Example:
|
||||
service = AutoTagService()
|
||||
status = service.get_models_status()
|
||||
|
||||
if status['gguf']['downloaded']:
|
||||
service.load_model('gguf')
|
||||
tags = service.generate_tags('/path/to/image.png')
|
||||
service.unload_model()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
models_dir: Optional[Path] = None,
|
||||
model_loader: Optional[ModelLoader] = None,
|
||||
):
|
||||
"""Initialize the AutoTag service.
|
||||
|
||||
Args:
|
||||
models_dir: Directory for storing models. If None, uses ComfyUI's
|
||||
folder_paths.models_dir / "LLM" path.
|
||||
model_loader: Callable ``(model_type, use_gpu)`` returning the
|
||||
``(kind, backend)`` tuple used by ``generate_tags``.
|
||||
Defaults to ``_load_model``, which imports the ML
|
||||
backends lazily. Tests inject a fake here.
|
||||
"""
|
||||
self.logger = get_logger("autotag.service")
|
||||
self._model_loader: ModelLoader = model_loader or self._load_model
|
||||
|
||||
# Determine models directory
|
||||
if models_dir:
|
||||
self.models_dir = Path(models_dir)
|
||||
else:
|
||||
# Use ComfyUI's folder_paths system
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
self.models_dir = Path(folder_paths.models_dir) / "LLM"
|
||||
except ImportError:
|
||||
# Fallback for standalone usage (not running in ComfyUI)
|
||||
self.logger.warning("folder_paths not available, using fallback path")
|
||||
self.models_dir = (
|
||||
Path(__file__).parent.parent.parent.parent / "models" / "LLM"
|
||||
)
|
||||
|
||||
self.logger.info(f"AutoTag service initialized. Models dir: {self.models_dir}")
|
||||
|
||||
# Current loaded tagger instance
|
||||
self._tagger = None
|
||||
self._current_model_type: Optional[str] = None
|
||||
self._custom_prompt: str = DEFAULT_PROMPT
|
||||
|
||||
# WD14 threshold defaults
|
||||
self._wd14_general_threshold: float = 0.35
|
||||
self._wd14_character_threshold: float = 0.85
|
||||
|
||||
@property
|
||||
def models_config(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get the models configuration dictionary."""
|
||||
return MODELS
|
||||
|
||||
@property
|
||||
def default_prompt(self) -> str:
|
||||
"""Get the default tag generation prompt."""
|
||||
return DEFAULT_PROMPT
|
||||
|
||||
@property
|
||||
def custom_prompt(self) -> str:
|
||||
"""Get the current custom prompt."""
|
||||
return self._custom_prompt
|
||||
|
||||
@custom_prompt.setter
|
||||
def custom_prompt(self, value: str):
|
||||
"""Set a custom prompt for tag generation."""
|
||||
self._custom_prompt = value
|
||||
|
||||
@property
|
||||
def wd14_general_threshold(self) -> float:
|
||||
return self._wd14_general_threshold
|
||||
|
||||
@wd14_general_threshold.setter
|
||||
def wd14_general_threshold(self, value: float):
|
||||
self._wd14_general_threshold = max(0.0, min(1.0, value))
|
||||
|
||||
@property
|
||||
def wd14_character_threshold(self) -> float:
|
||||
return self._wd14_character_threshold
|
||||
|
||||
@wd14_character_threshold.setter
|
||||
def wd14_character_threshold(self, value: float):
|
||||
self._wd14_character_threshold = max(0.0, min(1.0, value))
|
||||
|
||||
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get availability status for all model types.
|
||||
|
||||
Returns:
|
||||
Dictionary with status for each model type:
|
||||
{
|
||||
'gguf': {
|
||||
'name': 'GGUF (Recommended)',
|
||||
'description': '...',
|
||||
'size': '~4.5GB',
|
||||
'downloaded': True,
|
||||
'model_exists': True,
|
||||
'mmproj_exists': True, # GGUF only
|
||||
'model_path': '/path/to/model'
|
||||
},
|
||||
'hf': {...}
|
||||
}
|
||||
"""
|
||||
status = {}
|
||||
|
||||
for model_type, config in MODELS.items():
|
||||
model_status = {
|
||||
"name": config["name"],
|
||||
"description": config["description"],
|
||||
"size": config["size"],
|
||||
"downloaded": False,
|
||||
"model_path": None,
|
||||
}
|
||||
|
||||
if model_type == "gguf":
|
||||
model_exists, mmproj_exists = self._check_gguf_models()
|
||||
model_status["model_exists"] = model_exists
|
||||
model_status["mmproj_exists"] = mmproj_exists
|
||||
model_status["downloaded"] = model_exists and mmproj_exists
|
||||
if model_status["downloaded"]:
|
||||
model_status["model_path"] = str(
|
||||
self.models_dir / config["subdir"] / config["filename"]
|
||||
)
|
||||
elif model_type.startswith("wd14"):
|
||||
model_status["downloaded"] = self._check_wd14_model(model_type)
|
||||
if model_status["downloaded"]:
|
||||
model_status["model_path"] = str(
|
||||
self.models_dir / config["subdir"] / "model.onnx"
|
||||
)
|
||||
else: # hf
|
||||
model_status["downloaded"] = self._check_hf_model()
|
||||
if model_status["downloaded"]:
|
||||
# Get the actual path (local or cache)
|
||||
model_status["model_path"] = str(self._get_hf_model_path())
|
||||
|
||||
status[model_type] = model_status
|
||||
|
||||
return status
|
||||
|
||||
def _check_gguf_models(self) -> Tuple[bool, bool]:
|
||||
"""Check if GGUF model and mmproj files exist.
|
||||
|
||||
Returns:
|
||||
Tuple of (model_exists, mmproj_exists)
|
||||
"""
|
||||
config = MODELS["gguf"]
|
||||
gguf_dir = self.models_dir / config["subdir"]
|
||||
model_path = gguf_dir / config["filename"]
|
||||
mmproj_path = gguf_dir / config["mmproj_filename"]
|
||||
return model_path.exists(), mmproj_path.exists()
|
||||
|
||||
def _check_hf_model(self) -> bool:
|
||||
"""Check if HuggingFace model exists.
|
||||
|
||||
Checks both the local models directory and the HuggingFace cache.
|
||||
|
||||
Returns:
|
||||
True if model directory contains config.json
|
||||
"""
|
||||
config = MODELS["hf"]
|
||||
|
||||
# Check local directory first
|
||||
model_dir = self.models_dir / config["subdir"]
|
||||
self.logger.debug(f"Checking local HF model path: {model_dir}")
|
||||
if (model_dir / "config.json").exists():
|
||||
self.logger.debug("Found HF model in local directory")
|
||||
return True
|
||||
|
||||
# Check HuggingFace cache as fallback
|
||||
hf_cache_path = self._get_hf_cache_path(config["repo"])
|
||||
if hf_cache_path:
|
||||
self.logger.debug(f"Found HF model in cache: {hf_cache_path}")
|
||||
return True
|
||||
|
||||
self.logger.debug("HF model not found in local dir or cache")
|
||||
return False
|
||||
|
||||
def _get_hf_cache_path(self, repo_id: str) -> Optional[Path]:
|
||||
"""Get the path to a model in the HuggingFace cache.
|
||||
|
||||
Args:
|
||||
repo_id: The HuggingFace repo ID
|
||||
(e.g., 'fancyfeast/llama-joycaption-beta-one-hf-llava')
|
||||
|
||||
Returns:
|
||||
Path to the cached model directory, or None if not found
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
except ImportError:
|
||||
self.logger.debug("huggingface_hub not available for cache check")
|
||||
return None
|
||||
|
||||
try:
|
||||
cache_info = scan_cache_dir()
|
||||
for repo in cache_info.repos:
|
||||
if repo.repo_id == repo_id and repo.repo_type == "model":
|
||||
# Get the latest revision's snapshot path
|
||||
for revision in repo.revisions:
|
||||
snapshot_path = revision.snapshot_path
|
||||
if (Path(snapshot_path) / "config.json").exists():
|
||||
return Path(snapshot_path)
|
||||
except Exception as e:
|
||||
self.logger.debug(f"Error scanning HF cache: {e}")
|
||||
|
||||
return None
|
||||
|
||||
def _get_hf_model_path(self) -> Optional[Path]:
|
||||
"""Get the actual path to the HuggingFace model.
|
||||
|
||||
Checks local directory first, then HuggingFace cache.
|
||||
|
||||
Returns:
|
||||
Path to the model directory, or None if not found
|
||||
"""
|
||||
config = MODELS["hf"]
|
||||
|
||||
# Check local directory first
|
||||
model_dir = self.models_dir / config["subdir"]
|
||||
if (model_dir / "config.json").exists():
|
||||
return model_dir
|
||||
|
||||
# Check HuggingFace cache as fallback
|
||||
cache_path = self._get_hf_cache_path(config["repo"])
|
||||
if cache_path:
|
||||
return cache_path
|
||||
|
||||
return None
|
||||
|
||||
def _check_wd14_model(self, model_type: str) -> bool:
|
||||
"""Check if WD14 model files exist."""
|
||||
config = MODELS[model_type]
|
||||
model_dir = self.models_dir / config["subdir"]
|
||||
return all((model_dir / f).exists() for f in config["files"])
|
||||
|
||||
def _download_wd14_model(
|
||||
self,
|
||||
model_type: str,
|
||||
progress_callback: Optional[Callable[[str, float], None]] = None,
|
||||
) -> bool:
|
||||
"""Download WD14 model files from HuggingFace."""
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
config = MODELS[model_type]
|
||||
model_dir = self.models_dir / config["subdir"]
|
||||
model_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
files = config["files"]
|
||||
for i, filename in enumerate(files):
|
||||
file_path = model_dir / filename
|
||||
if not file_path.exists():
|
||||
progress = int((i / len(files)) * 80) + 10
|
||||
if progress_callback:
|
||||
progress_callback(f"Downloading {filename}...", progress)
|
||||
|
||||
self.logger.info(f"Downloading WD14 file: {filename}")
|
||||
hf_hub_download(
|
||||
repo_id=config["repo"],
|
||||
filename=filename,
|
||||
local_dir=str(model_dir),
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Download complete", 100)
|
||||
return True
|
||||
|
||||
def _load_wd14_tagger(self, model_type: str, use_gpu: bool = True):
|
||||
"""Load WD14 ONNX model and tag list."""
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
except ImportError:
|
||||
raise RuntimeError(
|
||||
"onnxruntime is required for WD14 models. "
|
||||
"Install with: pip install onnxruntime (or onnxruntime-gpu for CUDA)"
|
||||
)
|
||||
import csv
|
||||
|
||||
config = MODELS[model_type]
|
||||
model_dir = self.models_dir / config["subdir"]
|
||||
model_path = model_dir / "model.onnx"
|
||||
tags_path = model_dir / "selected_tags.csv"
|
||||
|
||||
self.logger.info(f"Loading WD14 model from {model_dir}...")
|
||||
|
||||
# Set up ONNX session
|
||||
providers = []
|
||||
if use_gpu:
|
||||
providers.append("CUDAExecutionProvider")
|
||||
providers.append("CPUExecutionProvider")
|
||||
|
||||
session = ort.InferenceSession(str(model_path), providers=providers)
|
||||
|
||||
# Read tag list from CSV
|
||||
tags_list = []
|
||||
with open(tags_path, "r", encoding="utf-8") as f:
|
||||
reader = csv.DictReader(f)
|
||||
for row in reader:
|
||||
tags_list.append(row)
|
||||
|
||||
self.logger.info(
|
||||
f"WD14 model loaded ({len(tags_list)} tags, "
|
||||
f"provider: {session.get_providers()[0]})"
|
||||
)
|
||||
return ("wd14", (session, tags_list))
|
||||
|
||||
def _generate_wd14(
|
||||
self,
|
||||
session,
|
||||
tags_list: List[Dict[str, str]],
|
||||
image: Image.Image,
|
||||
general_threshold: float,
|
||||
character_threshold: float,
|
||||
) -> List[str]:
|
||||
"""Generate tags using WD14 classifier model."""
|
||||
import numpy as np
|
||||
|
||||
# Get model input dimensions
|
||||
input_shape = session.get_inputs()[0].shape
|
||||
target_size = input_shape[2] # e.g. 448
|
||||
|
||||
# Pad to square with white background
|
||||
max_dim = max(image.size)
|
||||
padded = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
|
||||
padded.paste(
|
||||
image, ((max_dim - image.width) // 2, (max_dim - image.height) // 2)
|
||||
)
|
||||
|
||||
# Resize to model input size
|
||||
padded = padded.resize((target_size, target_size), Image.Resampling.BILINEAR)
|
||||
|
||||
# Convert to numpy: RGB -> BGR, float32, add batch dim
|
||||
img_array = np.array(padded, dtype=np.float32)
|
||||
img_array = img_array[:, :, ::-1] # RGB to BGR
|
||||
img_array = np.expand_dims(img_array, axis=0)
|
||||
|
||||
# Run inference
|
||||
input_name = session.get_inputs()[0].name
|
||||
output = session.run(None, {input_name: img_array})
|
||||
probabilities = output[0][0] # First output, first batch
|
||||
|
||||
# Filter tags by category and threshold
|
||||
result_tags = []
|
||||
for i, tag_info in enumerate(tags_list):
|
||||
if i >= len(probabilities):
|
||||
break
|
||||
prob = float(probabilities[i])
|
||||
category = tag_info.get("category", "0")
|
||||
tag_name = tag_info.get("name", "").strip()
|
||||
|
||||
if not tag_name:
|
||||
continue
|
||||
|
||||
# Category 0 = general, 4 = character, 9 = rating
|
||||
if category == "0" and prob >= general_threshold:
|
||||
result_tags.append((tag_name, prob))
|
||||
elif category == "4" and prob >= character_threshold:
|
||||
result_tags.append((tag_name, prob))
|
||||
|
||||
# Sort by confidence descending
|
||||
result_tags.sort(key=lambda x: x[1], reverse=True)
|
||||
return [tag for tag, _ in result_tags]
|
||||
|
||||
def download_model(
|
||||
self,
|
||||
model_type: str,
|
||||
progress_callback: Optional[Callable[[str, float], None]] = None,
|
||||
) -> bool:
|
||||
"""Download a model with optional progress updates.
|
||||
|
||||
Args:
|
||||
model_type: Either 'gguf' or 'hf'
|
||||
progress_callback: Optional callback(status_message, progress_percent)
|
||||
|
||||
Returns:
|
||||
True if download successful, False otherwise
|
||||
|
||||
Raises:
|
||||
ValueError: If model_type is not valid
|
||||
"""
|
||||
if model_type not in MODELS:
|
||||
raise ValueError(
|
||||
f"Invalid model type: {model_type}. "
|
||||
f"Must be one of: {', '.join(MODELS.keys())}"
|
||||
)
|
||||
|
||||
try:
|
||||
from huggingface_hub import ( # noqa: F401 - availability probe
|
||||
hf_hub_download,
|
||||
snapshot_download,
|
||||
)
|
||||
except ImportError:
|
||||
self.logger.error("huggingface_hub not installed")
|
||||
if progress_callback:
|
||||
progress_callback("Error: huggingface_hub not installed", 0)
|
||||
return False
|
||||
|
||||
try:
|
||||
if model_type == "gguf":
|
||||
return self._download_gguf_models(progress_callback)
|
||||
elif model_type.startswith("wd14"):
|
||||
return self._download_wd14_model(model_type, progress_callback)
|
||||
else:
|
||||
return self._download_hf_model(progress_callback)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Download failed: {e}")
|
||||
if progress_callback:
|
||||
progress_callback(f"Error: {str(e)}", 0)
|
||||
return False
|
||||
|
||||
def _download_gguf_models(
|
||||
self, progress_callback: Optional[Callable[[str, float], None]] = None
|
||||
) -> bool:
|
||||
"""Download GGUF model and mmproj files."""
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
config = MODELS["gguf"]
|
||||
gguf_dir = self.models_dir / config["subdir"]
|
||||
gguf_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model_path = gguf_dir / config["filename"]
|
||||
mmproj_path = gguf_dir / config["mmproj_filename"]
|
||||
|
||||
# Download main model
|
||||
if not model_path.exists():
|
||||
if progress_callback:
|
||||
progress_callback(f"Downloading {config['filename']}...", 10)
|
||||
|
||||
self.logger.info(f"Downloading GGUF model: {config['filename']}")
|
||||
hf_hub_download(
|
||||
repo_id=config["repo"],
|
||||
filename=config["filename"],
|
||||
local_dir=str(gguf_dir),
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
self.logger.info("GGUF model downloaded")
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Main model ready", 50)
|
||||
|
||||
# Download mmproj
|
||||
if not mmproj_path.exists():
|
||||
if progress_callback:
|
||||
progress_callback(f"Downloading {config['mmproj_filename']}...", 60)
|
||||
|
||||
self.logger.info(f"Downloading mmproj: {config['mmproj_filename']}")
|
||||
hf_hub_download(
|
||||
repo_id=config["mmproj_repo"],
|
||||
filename=config["mmproj_filename"],
|
||||
local_dir=str(gguf_dir),
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
self.logger.info("mmproj downloaded")
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Download complete", 100)
|
||||
|
||||
return True
|
||||
|
||||
def _download_hf_model(
|
||||
self, progress_callback: Optional[Callable[[str, float], None]] = None
|
||||
) -> bool:
|
||||
"""Download HuggingFace model."""
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
config = MODELS["hf"]
|
||||
model_dir = self.models_dir / config["subdir"]
|
||||
|
||||
if not self._check_hf_model():
|
||||
if progress_callback:
|
||||
progress_callback(f"Downloading {config['repo']}...", 10)
|
||||
|
||||
self.logger.info(f"Downloading HF model: {config['repo']}")
|
||||
snapshot_download(
|
||||
repo_id=config["repo"],
|
||||
local_dir=str(model_dir),
|
||||
local_dir_use_symlinks=False,
|
||||
)
|
||||
self.logger.info("HF model downloaded")
|
||||
|
||||
if progress_callback:
|
||||
progress_callback("Download complete", 100)
|
||||
|
||||
return True
|
||||
|
||||
def load_model(self, model_type: str, use_gpu: bool = True) -> bool:
|
||||
"""Load a model into memory for tag generation.
|
||||
|
||||
Args:
|
||||
model_type: Either 'gguf' or 'hf'
|
||||
use_gpu: Whether to use GPU acceleration (default True)
|
||||
|
||||
Returns:
|
||||
True if model loaded successfully
|
||||
|
||||
Raises:
|
||||
ValueError: If model_type is invalid
|
||||
RuntimeError: If model not downloaded or loading fails
|
||||
"""
|
||||
if model_type not in MODELS:
|
||||
raise ValueError(f"Invalid model type: {model_type}")
|
||||
|
||||
# Unload existing model first
|
||||
if self._tagger is not None:
|
||||
self.unload_model()
|
||||
|
||||
status = self.get_models_status()
|
||||
if not status[model_type]["downloaded"]:
|
||||
raise RuntimeError(f"Model {model_type} not downloaded")
|
||||
|
||||
try:
|
||||
self._tagger = self._model_loader(model_type, use_gpu)
|
||||
self._current_model_type = model_type
|
||||
self.logger.info(f"Model {model_type} loaded successfully")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to load model {model_type}: {e}")
|
||||
self._tagger = None
|
||||
self._current_model_type = None
|
||||
raise RuntimeError(f"Failed to load model: {e}")
|
||||
|
||||
def _load_model(self, model_type: str, use_gpu: bool = True) -> Tuple[str, Any]:
|
||||
"""Default loader: dispatch to a backend, importing it only now."""
|
||||
if model_type == "gguf":
|
||||
return self._load_gguf_tagger(use_gpu)
|
||||
if model_type.startswith("wd14"):
|
||||
return self._load_wd14_tagger(model_type, use_gpu)
|
||||
return self._load_hf_tagger()
|
||||
|
||||
def _load_gguf_tagger(self, use_gpu: bool = True):
|
||||
"""Load GGUF-based tagger."""
|
||||
from llama_cpp import Llama
|
||||
from llama_cpp.llama_chat_format import Llava15ChatHandler
|
||||
|
||||
config = MODELS["gguf"]
|
||||
gguf_dir = self.models_dir / config["subdir"]
|
||||
model_path = gguf_dir / config["filename"]
|
||||
mmproj_path = gguf_dir / config["mmproj_filename"]
|
||||
|
||||
self.logger.info("Loading GGUF model...")
|
||||
n_gpu_layers = -1 if use_gpu else 0
|
||||
|
||||
tagger = Llama(
|
||||
model_path=str(model_path),
|
||||
n_ctx=4096,
|
||||
n_batch=2048,
|
||||
n_threads=4,
|
||||
n_gpu_layers=n_gpu_layers,
|
||||
verbose=False,
|
||||
chat_handler=Llava15ChatHandler(clip_model_path=str(mmproj_path)),
|
||||
offload_kqv=True,
|
||||
)
|
||||
|
||||
self.logger.info("GGUF model loaded")
|
||||
return ("gguf", tagger)
|
||||
|
||||
def _load_hf_tagger(self, quantization: str = "8bit"):
|
||||
"""Load HuggingFace-based tagger."""
|
||||
import torch
|
||||
from transformers import (
|
||||
AutoProcessor,
|
||||
LlavaForConditionalGeneration,
|
||||
BitsAndBytesConfig,
|
||||
)
|
||||
|
||||
# Get the actual model path (local or cache)
|
||||
model_path = self._get_hf_model_path()
|
||||
if model_path is None:
|
||||
raise RuntimeError(
|
||||
"HuggingFace model not found in local directory or cache"
|
||||
)
|
||||
|
||||
self.logger.info(f"Loading HuggingFace model from {model_path}...")
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
processor = AutoProcessor.from_pretrained(str(model_path))
|
||||
model_kwargs = {"device_map": "cuda" if device == "cuda" else "cpu"}
|
||||
|
||||
if quantization == "8bit":
|
||||
qnt_config = BitsAndBytesConfig(
|
||||
load_in_8bit=True,
|
||||
bnb_8bit_compute_dtype=torch.float16,
|
||||
bnb_8bit_use_double_quant=True,
|
||||
llm_int8_skip_modules=["vision_tower", "multi_modal_projector"],
|
||||
)
|
||||
model = LlavaForConditionalGeneration.from_pretrained(
|
||||
str(model_path),
|
||||
torch_dtype=torch.float16,
|
||||
quantization_config=qnt_config,
|
||||
**model_kwargs,
|
||||
)
|
||||
else:
|
||||
model = LlavaForConditionalGeneration.from_pretrained(
|
||||
str(model_path), torch_dtype=torch.bfloat16, **model_kwargs
|
||||
)
|
||||
|
||||
model.eval()
|
||||
self.logger.info("HuggingFace model loaded")
|
||||
# Track the compute dtype for pixel_values conversion
|
||||
compute_dtype = torch.float16 if quantization == "8bit" else torch.bfloat16
|
||||
return ("hf", (model, processor, device, compute_dtype))
|
||||
|
||||
def unload_model(self):
|
||||
"""Unload the current model and free memory."""
|
||||
if self._tagger is not None:
|
||||
self.logger.info(f"Unloading model: {self._current_model_type}")
|
||||
del self._tagger
|
||||
self._tagger = None
|
||||
self._current_model_type = None
|
||||
gc.collect()
|
||||
|
||||
# Clear the CUDA cache only if torch is already loaded; importing
|
||||
# it here would pull in a heavy dependency the backend never used.
|
||||
torch = sys.modules.get("torch")
|
||||
if torch is not None and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
self.logger.info("Model unloaded, memory freed")
|
||||
|
||||
def is_model_loaded(self) -> bool:
|
||||
"""Check if a model is currently loaded."""
|
||||
return self._tagger is not None
|
||||
|
||||
def get_loaded_model_type(self) -> Optional[str]:
|
||||
"""Get the type of currently loaded model."""
|
||||
return self._current_model_type
|
||||
|
||||
def generate_tags(
|
||||
self,
|
||||
image_path: str,
|
||||
prompt: Optional[str] = None,
|
||||
general_threshold: Optional[float] = None,
|
||||
character_threshold: Optional[float] = None,
|
||||
) -> List[str]:
|
||||
"""Generate tags for an image.
|
||||
|
||||
Args:
|
||||
image_path: Path to the image file
|
||||
prompt: Custom prompt for tag generation (LLM models only).
|
||||
Uses default if None.
|
||||
general_threshold: Confidence threshold for general tags (WD14 only).
|
||||
character_threshold: Confidence threshold for character tags (WD14 only).
|
||||
|
||||
Returns:
|
||||
List of generated tags
|
||||
|
||||
Raises:
|
||||
RuntimeError: If no model is loaded
|
||||
FileNotFoundError: If image doesn't exist
|
||||
"""
|
||||
if self._tagger is None:
|
||||
raise RuntimeError("No model loaded. Call load_model() first.")
|
||||
|
||||
if not os.path.exists(image_path):
|
||||
raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path))
|
||||
|
||||
use_prompt = prompt or self._custom_prompt
|
||||
|
||||
# Load image
|
||||
image = Image.open(image_path)
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
|
||||
# Generate based on model type
|
||||
model_type, tagger_obj = self._tagger
|
||||
|
||||
if model_type == "gguf":
|
||||
raw_tags = self._generate_gguf(tagger_obj, image, use_prompt)
|
||||
return self._parse_tags(raw_tags)
|
||||
elif model_type == "wd14":
|
||||
session, tags_list = tagger_obj
|
||||
gen_thresh = (
|
||||
general_threshold
|
||||
if general_threshold is not None
|
||||
else self._wd14_general_threshold
|
||||
)
|
||||
char_thresh = (
|
||||
character_threshold
|
||||
if character_threshold is not None
|
||||
else self._wd14_character_threshold
|
||||
)
|
||||
return self._generate_wd14(
|
||||
session, tags_list, image, gen_thresh, char_thresh
|
||||
)
|
||||
else:
|
||||
model, processor, device, compute_dtype = tagger_obj
|
||||
raw_tags = self._generate_hf(
|
||||
model, processor, device, compute_dtype, image, use_prompt
|
||||
)
|
||||
return self._parse_tags(raw_tags)
|
||||
|
||||
def _generate_gguf(self, model, image: Image.Image, prompt: str) -> str:
|
||||
"""Generate tags using GGUF model."""
|
||||
import base64
|
||||
import io
|
||||
|
||||
# Resize image
|
||||
image = image.resize((336, 336), Image.Resampling.BILINEAR)
|
||||
|
||||
# Encode to base64
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
buffer.seek(0)
|
||||
img_base64 = base64.b64encode(buffer.read()).decode("utf-8")
|
||||
data_uri = f"data:image/png;base64,{img_base64}"
|
||||
|
||||
# Create message
|
||||
messages = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": data_uri}},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
# Generate
|
||||
response = model.create_chat_completion(
|
||||
messages=messages,
|
||||
max_tokens=512,
|
||||
temperature=0.6,
|
||||
top_p=0.9,
|
||||
stop=["</s>", "User:", "Assistant:"],
|
||||
stream=False,
|
||||
)
|
||||
|
||||
return response["choices"][0]["message"]["content"].strip()
|
||||
|
||||
def _generate_hf(
|
||||
self,
|
||||
model,
|
||||
processor,
|
||||
device: str,
|
||||
compute_dtype,
|
||||
image: Image.Image,
|
||||
prompt: str,
|
||||
) -> str:
|
||||
"""Generate tags using HuggingFace model."""
|
||||
import torch
|
||||
|
||||
# Resize image
|
||||
image = image.resize((336, 336), Image.Resampling.LANCZOS)
|
||||
|
||||
convo = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
]
|
||||
convo_string = processor.apply_chat_template(
|
||||
convo, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(
|
||||
device
|
||||
)
|
||||
|
||||
# Convert pixel_values to the model's compute dtype
|
||||
# (float16 for 8-bit quantized, bfloat16 for non-quantized)
|
||||
if "pixel_values" in inputs and inputs["pixel_values"] is not None:
|
||||
inputs["pixel_values"] = inputs["pixel_values"].to(compute_dtype)
|
||||
|
||||
with torch.inference_mode(), torch.cuda.amp.autocast(enabled=True):
|
||||
generate_ids = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=512,
|
||||
do_sample=True,
|
||||
temperature=0.6,
|
||||
top_p=0.9,
|
||||
use_cache=True,
|
||||
)[0]
|
||||
|
||||
generate_ids = generate_ids[inputs["input_ids"].shape[1] :]
|
||||
return processor.tokenizer.decode(
|
||||
generate_ids, skip_special_tokens=True
|
||||
).strip()
|
||||
|
||||
def _parse_tags(self, raw_output: str) -> List[str]:
|
||||
"""Parse raw model output into a list of clean tags.
|
||||
|
||||
Args:
|
||||
raw_output: Raw text output from the model
|
||||
|
||||
Returns:
|
||||
List of cleaned, deduplicated tags
|
||||
"""
|
||||
# Patterns to exclude
|
||||
exclude_prefixes = (
|
||||
"copyright:",
|
||||
"meta:",
|
||||
"photo_",
|
||||
"photo:",
|
||||
"prepend:",
|
||||
"append:",
|
||||
)
|
||||
|
||||
# Split by common delimiters
|
||||
tags = []
|
||||
|
||||
# Handle comma-separated tags
|
||||
for part in raw_output.split(","):
|
||||
tag = part.strip().lower()
|
||||
# Remove any quotes or extra characters
|
||||
tag = tag.strip("\"'")
|
||||
# Replace spaces with underscores (Danbooru style)
|
||||
tag = tag.replace(" ", "_")
|
||||
# Remove empty tags
|
||||
if tag and len(tag) > 1:
|
||||
# Filter out unwanted tag patterns
|
||||
if not tag.startswith(exclude_prefixes):
|
||||
tags.append(tag)
|
||||
|
||||
# Deduplicate while preserving order
|
||||
seen = set()
|
||||
unique_tags = []
|
||||
for tag in tags:
|
||||
if tag not in seen:
|
||||
seen.add(tag)
|
||||
unique_tags.append(tag)
|
||||
|
||||
return unique_tags
|
||||
|
||||
|
||||
# Singleton instance for API use
|
||||
_service_instance: Optional[AutoTagService] = None
|
||||
|
||||
|
||||
def get_autotag_service() -> AutoTagService:
|
||||
"""Get or create the singleton AutoTagService instance."""
|
||||
global _service_instance
|
||||
if _service_instance is None:
|
||||
_service_instance = AutoTagService()
|
||||
return _service_instance
|
||||
@@ -22,39 +22,144 @@ Example:
|
||||
extension_name = "PromptManager"
|
||||
|
||||
# Get server instance and routes (same pattern as ComfyUI_Assets)
|
||||
from server import PromptServer
|
||||
from server import PromptServer # noqa: E402
|
||||
|
||||
server_instance = PromptServer.instance
|
||||
routes = server_instance.routes
|
||||
|
||||
# Extension info
|
||||
extension_uri = None # Will be set in __init__.py
|
||||
|
||||
import os
|
||||
from typing import Dict, Any, List
|
||||
import os # noqa: E402
|
||||
from pathlib import Path # noqa: E402
|
||||
from typing import Any, Dict, List, Optional, Tuple # noqa: E402
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, current_dir)
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
# Initialize logger for config operations
|
||||
config_logger = get_logger('prompt_manager.config')
|
||||
config_logger = get_logger("prompt_manager.config")
|
||||
|
||||
# Environment variable listing extra directories (os.pathsep-separated) that
|
||||
# may be used as gallery roots in addition to ComfyUI's own directories.
|
||||
EXTRA_GALLERY_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
|
||||
|
||||
# Environment variable overriding where config.json is read from and saved to.
|
||||
CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH"
|
||||
|
||||
# ComfyUI folder_paths getters whose directories are valid gallery parents.
|
||||
_COMFYUI_DIRECTORY_GETTERS = (
|
||||
"get_output_directory",
|
||||
"get_input_directory",
|
||||
"get_temp_directory",
|
||||
"get_user_directory",
|
||||
)
|
||||
|
||||
|
||||
def _canonical_path(path: str) -> str:
|
||||
"""Return a path normalised for comparison (realpath + normcase)."""
|
||||
return os.path.normcase(os.path.realpath(path))
|
||||
|
||||
|
||||
def write_private_json(path: str, data: Any) -> None:
|
||||
"""Write ``data`` as JSON to ``path`` readable by the current user only.
|
||||
|
||||
config.json can hold an API key (integrations.lora_manager), so it is
|
||||
created with mode 0600 and, on POSIX, an existing more open file is
|
||||
tightened too. Windows has no POSIX modes; there the chmod is a best
|
||||
effort that only affects the read-only bit. Parent directories are
|
||||
created as needed. Raises OSError when the file cannot be written.
|
||||
"""
|
||||
import json
|
||||
|
||||
parent_dir = os.path.dirname(path)
|
||||
if parent_dir:
|
||||
os.makedirs(parent_dir, exist_ok=True)
|
||||
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
if os.name == "posix":
|
||||
os.chmod(path, 0o600)
|
||||
|
||||
|
||||
def _is_filesystem_root(path: str) -> bool:
|
||||
"""True for '/' on POSIX and drive roots such as 'C:\\' on Windows."""
|
||||
return os.path.dirname(path) == path
|
||||
|
||||
|
||||
def _import_folder_paths():
|
||||
"""ComfyUI's folder_paths module, or None outside ComfyUI."""
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
return None
|
||||
return folder_paths
|
||||
|
||||
|
||||
def _comfyui_base_directory() -> Optional[str]:
|
||||
"""ComfyUI's base directory: folder_paths.base_path, else the output dir's
|
||||
parent."""
|
||||
folder_paths = _import_folder_paths()
|
||||
if folder_paths is None:
|
||||
return None
|
||||
base = getattr(folder_paths, "base_path", None)
|
||||
if isinstance(base, str) and base:
|
||||
return base
|
||||
getter = getattr(folder_paths, "get_output_directory", None)
|
||||
if getter is None:
|
||||
return None
|
||||
try:
|
||||
output_dir = getter()
|
||||
except Exception:
|
||||
return None
|
||||
if isinstance(output_dir, str) and output_dir:
|
||||
return os.path.dirname(os.path.realpath(output_dir))
|
||||
return None
|
||||
|
||||
|
||||
def _comfyui_directories() -> List[str]:
|
||||
"""Directories reported by ComfyUI's folder_paths module, if importable."""
|
||||
folder_paths = _import_folder_paths()
|
||||
if folder_paths is None:
|
||||
return []
|
||||
|
||||
found = []
|
||||
for getter_name in _COMFYUI_DIRECTORY_GETTERS:
|
||||
getter = getattr(folder_paths, getter_name, None)
|
||||
if getter is None:
|
||||
continue
|
||||
try:
|
||||
directory = getter()
|
||||
except Exception:
|
||||
continue
|
||||
if isinstance(directory, str) and directory:
|
||||
found.append(directory)
|
||||
return found
|
||||
|
||||
|
||||
def _extra_gallery_roots() -> List[str]:
|
||||
"""Directories the user opted in through PROMPT_MANAGER_EXTRA_GALLERY_ROOTS."""
|
||||
raw = os.environ.get(EXTRA_GALLERY_ROOTS_ENV, "")
|
||||
return [entry.strip() for entry in raw.split(os.pathsep) if entry.strip()]
|
||||
|
||||
|
||||
class GalleryConfig:
|
||||
"""Configuration class for the gallery monitoring and image processing system.
|
||||
|
||||
|
||||
This class manages all settings related to automatic image monitoring,
|
||||
prompt tracking, database cleanup, web interface display, and performance
|
||||
optimization for the gallery functionality.
|
||||
|
||||
|
||||
All configuration values are class attributes that can be modified at runtime
|
||||
or loaded from external configuration files.
|
||||
|
||||
|
||||
Attributes:
|
||||
MONITORING_ENABLED (bool): Enable/disable automatic image monitoring
|
||||
MONITORING_DIRECTORIES (List[str]): Directories to monitor for new images
|
||||
@@ -71,86 +176,89 @@ class GalleryConfig:
|
||||
MAX_CONCURRENT_PROCESSING (int): Maximum concurrent image processing tasks
|
||||
METADATA_EXTRACTION_TIMEOUT (int): Timeout for metadata extraction operations
|
||||
"""
|
||||
|
||||
|
||||
# Image monitoring settings
|
||||
MONITORING_ENABLED = True
|
||||
MONITORING_DIRECTORIES = [] # Auto-detect if empty
|
||||
SUPPORTED_EXTENSIONS = ['.png', '.jpg', '.jpeg', '.webp', '.gif']
|
||||
SUPPORTED_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".gif"]
|
||||
PROCESSING_DELAY = 2.0 # Seconds to wait before processing new files
|
||||
|
||||
|
||||
# Prompt tracking settings
|
||||
PROMPT_TIMEOUT = 120 # Seconds to keep prompt context active
|
||||
PROMPT_TIMEOUT = (
|
||||
600 # Seconds to keep prompt context active (10 min for long generations)
|
||||
)
|
||||
CLEANUP_INTERVAL = 300 # Seconds between cleanup of expired prompts
|
||||
|
||||
|
||||
# Database settings
|
||||
AUTO_CLEANUP_MISSING_FILES = True
|
||||
MAX_IMAGE_AGE_DAYS = 365 # Clean up images older than this
|
||||
|
||||
|
||||
# Web interface settings
|
||||
IMAGES_PER_PAGE = 20
|
||||
THUMBNAIL_SIZE = 256
|
||||
ENABLE_SEARCH = True
|
||||
ENABLE_METADATA_VIEW = True
|
||||
|
||||
|
||||
# Performance settings
|
||||
MAX_CONCURRENT_PROCESSING = 3
|
||||
METADATA_EXTRACTION_TIMEOUT = 10 # Seconds
|
||||
|
||||
|
||||
@classmethod
|
||||
def get_config(cls) -> Dict[str, Any]:
|
||||
"""Get the complete gallery configuration as a structured dictionary.
|
||||
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: A nested dictionary containing all gallery configuration
|
||||
sections: monitoring, tracking, database, web_interface, and performance.
|
||||
Each section contains the relevant configuration parameters as key-value pairs.
|
||||
|
||||
sections: monitoring, tracking, database, web_interface, and
|
||||
performance. Each section contains the relevant configuration
|
||||
parameters as key-value pairs.
|
||||
|
||||
Example:
|
||||
config = GalleryConfig.get_config()
|
||||
monitoring_enabled = config['monitoring']['enabled']
|
||||
images_per_page = config['web_interface']['images_per_page']
|
||||
"""
|
||||
return {
|
||||
'monitoring': {
|
||||
'enabled': cls.MONITORING_ENABLED,
|
||||
'directories': cls.MONITORING_DIRECTORIES,
|
||||
'extensions': cls.SUPPORTED_EXTENSIONS,
|
||||
'processing_delay': cls.PROCESSING_DELAY
|
||||
"monitoring": {
|
||||
"enabled": cls.MONITORING_ENABLED,
|
||||
"directories": cls.MONITORING_DIRECTORIES,
|
||||
"extensions": cls.SUPPORTED_EXTENSIONS,
|
||||
"processing_delay": cls.PROCESSING_DELAY,
|
||||
},
|
||||
'tracking': {
|
||||
'prompt_timeout': cls.PROMPT_TIMEOUT,
|
||||
'cleanup_interval': cls.CLEANUP_INTERVAL
|
||||
"tracking": {
|
||||
"prompt_timeout": cls.PROMPT_TIMEOUT,
|
||||
"cleanup_interval": cls.CLEANUP_INTERVAL,
|
||||
},
|
||||
'database': {
|
||||
'auto_cleanup': cls.AUTO_CLEANUP_MISSING_FILES,
|
||||
'max_image_age_days': cls.MAX_IMAGE_AGE_DAYS
|
||||
"database": {
|
||||
"auto_cleanup": cls.AUTO_CLEANUP_MISSING_FILES,
|
||||
"max_image_age_days": cls.MAX_IMAGE_AGE_DAYS,
|
||||
},
|
||||
'web_interface': {
|
||||
'images_per_page': cls.IMAGES_PER_PAGE,
|
||||
'thumbnail_size': cls.THUMBNAIL_SIZE,
|
||||
'enable_search': cls.ENABLE_SEARCH,
|
||||
'enable_metadata_view': cls.ENABLE_METADATA_VIEW
|
||||
"web_interface": {
|
||||
"images_per_page": cls.IMAGES_PER_PAGE,
|
||||
"thumbnail_size": cls.THUMBNAIL_SIZE,
|
||||
"enable_search": cls.ENABLE_SEARCH,
|
||||
"enable_metadata_view": cls.ENABLE_METADATA_VIEW,
|
||||
},
|
||||
"performance": {
|
||||
"max_concurrent_processing": cls.MAX_CONCURRENT_PROCESSING,
|
||||
"metadata_extraction_timeout": cls.METADATA_EXTRACTION_TIMEOUT,
|
||||
},
|
||||
'performance': {
|
||||
'max_concurrent_processing': cls.MAX_CONCURRENT_PROCESSING,
|
||||
'metadata_extraction_timeout': cls.METADATA_EXTRACTION_TIMEOUT
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@classmethod
|
||||
def update_config(cls, new_config: Dict[str, Any]):
|
||||
"""Update gallery configuration attributes from a dictionary.
|
||||
|
||||
|
||||
Takes a nested dictionary with gallery configuration sections and updates
|
||||
the corresponding class attributes. Only updates attributes that are
|
||||
present in the input dictionary, leaving others unchanged.
|
||||
|
||||
|
||||
Args:
|
||||
new_config (Dict[str, Any]): Nested dictionary containing gallery
|
||||
configuration updates. Should follow the same structure as returned
|
||||
by get_config(). Valid top-level keys are: 'monitoring', 'tracking',
|
||||
'database', 'web_interface', 'performance'.
|
||||
|
||||
|
||||
Example:
|
||||
gallery_settings = {
|
||||
'monitoring': {'enabled': False},
|
||||
@@ -158,58 +266,218 @@ class GalleryConfig:
|
||||
}
|
||||
GalleryConfig.update_config(gallery_settings)
|
||||
"""
|
||||
monitoring = new_config.get('monitoring', {})
|
||||
if 'enabled' in monitoring:
|
||||
cls.MONITORING_ENABLED = monitoring['enabled']
|
||||
if 'directories' in monitoring:
|
||||
cls.MONITORING_DIRECTORIES = monitoring['directories']
|
||||
if 'extensions' in monitoring:
|
||||
cls.SUPPORTED_EXTENSIONS = monitoring['extensions']
|
||||
if 'processing_delay' in monitoring:
|
||||
cls.PROCESSING_DELAY = monitoring['processing_delay']
|
||||
|
||||
tracking = new_config.get('tracking', {})
|
||||
if 'prompt_timeout' in tracking:
|
||||
cls.PROMPT_TIMEOUT = tracking['prompt_timeout']
|
||||
if 'cleanup_interval' in tracking:
|
||||
cls.CLEANUP_INTERVAL = tracking['cleanup_interval']
|
||||
|
||||
database = new_config.get('database', {})
|
||||
if 'auto_cleanup' in database:
|
||||
cls.AUTO_CLEANUP_MISSING_FILES = database['auto_cleanup']
|
||||
if 'max_image_age_days' in database:
|
||||
cls.MAX_IMAGE_AGE_DAYS = database['max_image_age_days']
|
||||
|
||||
web_interface = new_config.get('web_interface', {})
|
||||
if 'images_per_page' in web_interface:
|
||||
cls.IMAGES_PER_PAGE = web_interface['images_per_page']
|
||||
if 'thumbnail_size' in web_interface:
|
||||
cls.THUMBNAIL_SIZE = web_interface['thumbnail_size']
|
||||
if 'enable_search' in web_interface:
|
||||
cls.ENABLE_SEARCH = web_interface['enable_search']
|
||||
if 'enable_metadata_view' in web_interface:
|
||||
cls.ENABLE_METADATA_VIEW = web_interface['enable_metadata_view']
|
||||
|
||||
performance = new_config.get('performance', {})
|
||||
if 'max_concurrent_processing' in performance:
|
||||
cls.MAX_CONCURRENT_PROCESSING = performance['max_concurrent_processing']
|
||||
if 'metadata_extraction_timeout' in performance:
|
||||
cls.METADATA_EXTRACTION_TIMEOUT = performance['metadata_extraction_timeout']
|
||||
monitoring = new_config.get("monitoring", {})
|
||||
if "enabled" in monitoring:
|
||||
cls.MONITORING_ENABLED = monitoring["enabled"]
|
||||
if "directories" in monitoring:
|
||||
directories = cls._validated_directories(monitoring["directories"])
|
||||
if directories is not None:
|
||||
cls.MONITORING_DIRECTORIES = directories
|
||||
if "extensions" in monitoring:
|
||||
cls.SUPPORTED_EXTENSIONS = monitoring["extensions"]
|
||||
if "processing_delay" in monitoring:
|
||||
cls.PROCESSING_DELAY = monitoring["processing_delay"]
|
||||
|
||||
tracking = new_config.get("tracking", {})
|
||||
if "prompt_timeout" in tracking:
|
||||
cls.PROMPT_TIMEOUT = tracking["prompt_timeout"]
|
||||
if "cleanup_interval" in tracking:
|
||||
cls.CLEANUP_INTERVAL = tracking["cleanup_interval"]
|
||||
|
||||
database = new_config.get("database", {})
|
||||
if "auto_cleanup" in database:
|
||||
cls.AUTO_CLEANUP_MISSING_FILES = database["auto_cleanup"]
|
||||
if "max_image_age_days" in database:
|
||||
cls.MAX_IMAGE_AGE_DAYS = database["max_image_age_days"]
|
||||
|
||||
web_interface = new_config.get("web_interface", {})
|
||||
if "images_per_page" in web_interface:
|
||||
cls.IMAGES_PER_PAGE = web_interface["images_per_page"]
|
||||
if "thumbnail_size" in web_interface:
|
||||
cls.THUMBNAIL_SIZE = web_interface["thumbnail_size"]
|
||||
if "enable_search" in web_interface:
|
||||
cls.ENABLE_SEARCH = web_interface["enable_search"]
|
||||
if "enable_metadata_view" in web_interface:
|
||||
cls.ENABLE_METADATA_VIEW = web_interface["enable_metadata_view"]
|
||||
|
||||
performance = new_config.get("performance", {})
|
||||
if "max_concurrent_processing" in performance:
|
||||
cls.MAX_CONCURRENT_PROCESSING = performance["max_concurrent_processing"]
|
||||
if "metadata_extraction_timeout" in performance:
|
||||
cls.METADATA_EXTRACTION_TIMEOUT = performance["metadata_extraction_timeout"]
|
||||
|
||||
@classmethod
|
||||
def _validated_directories(cls, directories: Any) -> Optional[List[str]]:
|
||||
"""Gallery roots from a config payload that pass validate_gallery_root.
|
||||
|
||||
A hand-edited config.json is as untrusted as a settings request:
|
||||
entries that fail validation are dropped with a warning (naming only
|
||||
the basename), and a value that is not a list is ignored entirely
|
||||
(``None`` is returned so the current roots stay as they are).
|
||||
"""
|
||||
if not isinstance(directories, list):
|
||||
config_logger.warning(
|
||||
"Ignoring gallery monitoring directories: expected a list"
|
||||
)
|
||||
return None
|
||||
kept = []
|
||||
for entry in directories:
|
||||
ok, reason = cls.validate_gallery_root(entry)
|
||||
if ok:
|
||||
kept.append(entry)
|
||||
continue
|
||||
label = os.path.basename(entry) if isinstance(entry, str) else entry
|
||||
config_logger.warning(f"Ignoring gallery root {label!r}: {reason}")
|
||||
return kept
|
||||
|
||||
@classmethod
|
||||
def allowed_gallery_parents(cls) -> List[str]:
|
||||
"""Canonical directories under which gallery roots may live.
|
||||
|
||||
Filesystem roots are never allowed as parents, even when listed in
|
||||
the environment, because that would re-open the whole disk.
|
||||
"""
|
||||
parents = []
|
||||
for candidate in _comfyui_directories() + _extra_gallery_roots():
|
||||
try:
|
||||
canonical = _canonical_path(candidate)
|
||||
except (OSError, ValueError):
|
||||
continue
|
||||
if _is_filesystem_root(canonical) or not os.path.isdir(canonical):
|
||||
continue
|
||||
if canonical not in parents:
|
||||
parents.append(canonical)
|
||||
return parents
|
||||
|
||||
@classmethod
|
||||
def path_anchors(cls) -> List[str]:
|
||||
"""Canonical directories that relative gallery paths are resolved against.
|
||||
|
||||
The ComfyUI base directory comes first, followed by the parent of each
|
||||
directory listed in PROMPT_MANAGER_EXTRA_GALLERY_ROOTS, so an extra
|
||||
root is addressed by its own name (``gallery/sub``) rather than by an
|
||||
absolute path. The same anchors drive the public (relative) form of
|
||||
paths in API responses.
|
||||
"""
|
||||
anchors = []
|
||||
base = _comfyui_base_directory()
|
||||
candidates = [base] if base else []
|
||||
candidates.extend(os.path.dirname(root) for root in _extra_gallery_roots())
|
||||
for candidate in candidates:
|
||||
try:
|
||||
canonical = _canonical_path(candidate)
|
||||
except (OSError, ValueError):
|
||||
continue
|
||||
if canonical not in anchors:
|
||||
anchors.append(canonical)
|
||||
return anchors
|
||||
|
||||
@classmethod
|
||||
def resolve_gallery_root(cls, path: str) -> str:
|
||||
"""Canonical absolute path for a gallery root given in any accepted form.
|
||||
|
||||
Absolute paths are canonicalised as-is. Relative paths are tried
|
||||
against each :meth:`path_anchors` entry and the first existing match
|
||||
wins; otherwise the first anchor (or the current directory when there
|
||||
is none) is used, so the caller's existence check reports it.
|
||||
"""
|
||||
path = path.strip()
|
||||
if os.path.isabs(path):
|
||||
return _canonical_path(path)
|
||||
anchors = cls.path_anchors()
|
||||
for anchor in anchors:
|
||||
candidate = os.path.join(anchor, path)
|
||||
if os.path.exists(candidate):
|
||||
return _canonical_path(candidate)
|
||||
return _canonical_path(os.path.join(anchors[0], path) if anchors else path)
|
||||
|
||||
@classmethod
|
||||
def validate_gallery_root(cls, path: Any) -> Tuple[bool, str]:
|
||||
"""Check whether ``path`` may be used as a gallery root.
|
||||
|
||||
Returns:
|
||||
(True, "") when the directory lies inside one of
|
||||
:meth:`allowed_gallery_parents`; otherwise (False, reason).
|
||||
"""
|
||||
if not isinstance(path, str) or not path.strip():
|
||||
return False, "Gallery root must be a non-empty path"
|
||||
|
||||
try:
|
||||
canonical = cls.resolve_gallery_root(path)
|
||||
except (OSError, ValueError):
|
||||
return False, "Gallery root could not be resolved"
|
||||
|
||||
if _is_filesystem_root(canonical):
|
||||
return False, "Gallery root cannot be a filesystem root"
|
||||
if canonical == _canonical_path(os.path.expanduser("~")):
|
||||
return False, "Gallery root cannot be the home directory"
|
||||
if not os.path.exists(canonical):
|
||||
return False, "Gallery root does not exist"
|
||||
if not os.path.isdir(canonical):
|
||||
return False, "Gallery root is not a directory"
|
||||
|
||||
candidate = Path(canonical)
|
||||
for parent in cls.allowed_gallery_parents():
|
||||
if candidate.is_relative_to(Path(parent)):
|
||||
return True, ""
|
||||
|
||||
return (
|
||||
False,
|
||||
"Gallery root must be inside a ComfyUI directory "
|
||||
f"(output, input, temp, user) or one listed in {EXTRA_GALLERY_ROOTS_ENV}",
|
||||
)
|
||||
|
||||
|
||||
class IntegrationConfig:
|
||||
"""Configuration for third-party extension integrations.
|
||||
|
||||
Manages opt-in integration settings for extensions like LoraManager.
|
||||
All integrations are disabled by default so PromptManager works standalone.
|
||||
"""
|
||||
|
||||
# LoraManager integration
|
||||
LORA_MANAGER_ENABLED = False
|
||||
LORA_MANAGER_PATH = "" # Auto-detected if empty
|
||||
LORA_TRIGGER_WORDS_ENABLED = False # Auto-inject trigger words into prompts
|
||||
CIVITAI_API_KEY = "" # Required to download NSFW example images
|
||||
|
||||
@classmethod
|
||||
def get_config(cls) -> Dict[str, Any]:
|
||||
return {
|
||||
"lora_manager": {
|
||||
"enabled": cls.LORA_MANAGER_ENABLED,
|
||||
"path": cls.LORA_MANAGER_PATH,
|
||||
"trigger_words_enabled": cls.LORA_TRIGGER_WORDS_ENABLED,
|
||||
"civitai_api_key": cls.CIVITAI_API_KEY,
|
||||
},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def update_config(cls, new_config: Dict[str, Any]):
|
||||
lora = new_config.get("lora_manager", {})
|
||||
if "enabled" in lora:
|
||||
cls.LORA_MANAGER_ENABLED = lora["enabled"]
|
||||
if "path" in lora:
|
||||
cls.LORA_MANAGER_PATH = lora["path"]
|
||||
if "trigger_words_enabled" in lora:
|
||||
cls.LORA_TRIGGER_WORDS_ENABLED = lora["trigger_words_enabled"]
|
||||
if "civitai_api_key" in lora:
|
||||
cls.CIVITAI_API_KEY = lora["civitai_api_key"]
|
||||
|
||||
|
||||
class PromptManagerConfig:
|
||||
"""Main configuration class for PromptManager core functionality.
|
||||
|
||||
|
||||
This class manages configuration for database operations, web UI behavior,
|
||||
performance settings, and integrates gallery configuration. It provides
|
||||
methods for loading and saving configuration from/to JSON files.
|
||||
|
||||
|
||||
The configuration is organized into logical sections:
|
||||
- Database: Settings for SQLite operations and data management
|
||||
- Web UI: User interface behavior and display options
|
||||
- Performance: Optimization and resource management settings
|
||||
- Gallery: Embedded gallery configuration (via GalleryConfig)
|
||||
|
||||
|
||||
Attributes:
|
||||
DEFAULT_DB_PATH (str): Default path for the SQLite database file
|
||||
ENABLE_DUPLICATE_DETECTION (bool): Enable automatic duplicate detection
|
||||
@@ -221,82 +489,127 @@ class PromptManagerConfig:
|
||||
ENABLE_FUZZY_SEARCH (bool): Enable fuzzy search capabilities
|
||||
AUTO_BACKUP_INTERVAL (int): Hours between automatic database backups
|
||||
"""
|
||||
|
||||
|
||||
# Database settings
|
||||
DEFAULT_DB_PATH = "prompts.db"
|
||||
ENABLE_DUPLICATE_DETECTION = True
|
||||
ENABLE_AUTO_SAVE = True
|
||||
|
||||
|
||||
# Web UI settings
|
||||
RESULT_TIMEOUT = 5 # Seconds to auto-hide results in ComfyUI node
|
||||
SHOW_TEST_BUTTON = False # Show API test button in node UI
|
||||
WEBUI_DISPLAY_MODE = 'newtab' # 'popup' or 'newtab'
|
||||
|
||||
WEBUI_DISPLAY_MODE = "newtab" # 'popup' or 'newtab'
|
||||
INFINITE_SCROLL = False # Dashboard list: load the next page on scroll by default
|
||||
|
||||
# Performance settings
|
||||
MAX_SEARCH_RESULTS = 100
|
||||
ENABLE_FUZZY_SEARCH = False # Requires fuzzywuzzy
|
||||
AUTO_BACKUP_INTERVAL = 24 # Hours
|
||||
|
||||
# Threads for the output scan and thumbnail generation; half the cores by
|
||||
# default so ComfyUI's own generation keeps CPU headroom. 1..max_worker_threads().
|
||||
DEFAULT_WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2)
|
||||
WORKER_THREADS = DEFAULT_WORKER_THREADS
|
||||
|
||||
@classmethod
|
||||
def max_worker_threads(cls) -> int:
|
||||
"""Largest accepted worker-thread count: the detected core count."""
|
||||
return os.cpu_count() or 1
|
||||
|
||||
@classmethod
|
||||
def clamp_worker_threads(cls, value, fallback: int) -> int:
|
||||
"""Coerce a worker count from config.json into ``[1, max_worker_threads()]``.
|
||||
|
||||
Non-numeric values (and booleans) return ``fallback`` so a hand-edited
|
||||
file cannot disable the setting.
|
||||
"""
|
||||
if isinstance(value, bool):
|
||||
return fallback
|
||||
if isinstance(value, str) and value.strip().isdigit():
|
||||
value = int(value.strip())
|
||||
if not isinstance(value, int):
|
||||
return fallback
|
||||
return max(1, min(value, cls.max_worker_threads()))
|
||||
|
||||
@classmethod
|
||||
def get_config(cls) -> Dict[str, Any]:
|
||||
"""Get the complete PromptManager configuration as a structured dictionary.
|
||||
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: A nested dictionary containing all configuration sections:
|
||||
- database: Database-related settings
|
||||
- web_ui: Web interface configuration
|
||||
- performance: Performance and optimization settings
|
||||
- gallery: Complete gallery configuration (from GalleryConfig)
|
||||
|
||||
|
||||
Example:
|
||||
config = PromptManagerConfig.get_config()
|
||||
db_path = config['database']['default_path']
|
||||
max_results = config['performance']['max_search_results']
|
||||
"""
|
||||
return {
|
||||
'database': {
|
||||
'default_path': cls.DEFAULT_DB_PATH,
|
||||
'enable_duplicate_detection': cls.ENABLE_DUPLICATE_DETECTION,
|
||||
'enable_auto_save': cls.ENABLE_AUTO_SAVE
|
||||
"database": {
|
||||
"default_path": cls.DEFAULT_DB_PATH,
|
||||
"enable_duplicate_detection": cls.ENABLE_DUPLICATE_DETECTION,
|
||||
"enable_auto_save": cls.ENABLE_AUTO_SAVE,
|
||||
},
|
||||
'web_ui': {
|
||||
'result_timeout': cls.RESULT_TIMEOUT,
|
||||
'show_test_button': cls.SHOW_TEST_BUTTON,
|
||||
'webui_display_mode': cls.WEBUI_DISPLAY_MODE
|
||||
"web_ui": {
|
||||
"result_timeout": cls.RESULT_TIMEOUT,
|
||||
"show_test_button": cls.SHOW_TEST_BUTTON,
|
||||
"webui_display_mode": cls.WEBUI_DISPLAY_MODE,
|
||||
"infinite_scroll": cls.INFINITE_SCROLL,
|
||||
},
|
||||
'performance': {
|
||||
'max_search_results': cls.MAX_SEARCH_RESULTS,
|
||||
'enable_fuzzy_search': cls.ENABLE_FUZZY_SEARCH,
|
||||
'auto_backup_interval': cls.AUTO_BACKUP_INTERVAL
|
||||
"performance": {
|
||||
"max_search_results": cls.MAX_SEARCH_RESULTS,
|
||||
"enable_fuzzy_search": cls.ENABLE_FUZZY_SEARCH,
|
||||
"auto_backup_interval": cls.AUTO_BACKUP_INTERVAL,
|
||||
"worker_threads": cls.WORKER_THREADS,
|
||||
},
|
||||
'gallery': GalleryConfig.get_config()
|
||||
"gallery": GalleryConfig.get_config(),
|
||||
"integrations": IntegrationConfig.get_config(),
|
||||
}
|
||||
|
||||
|
||||
@classmethod
|
||||
def load_from_file(cls, config_path: str):
|
||||
def get_config_path(cls) -> str:
|
||||
"""Path of the persisted config.json.
|
||||
|
||||
Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file
|
||||
lives at the repository root next to ``pyproject.toml``.
|
||||
"""
|
||||
override = os.environ.get(CONFIG_PATH_ENV, "").strip()
|
||||
if override:
|
||||
return override
|
||||
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
return os.path.join(repo_root, "config.json")
|
||||
|
||||
@classmethod
|
||||
def load_from_file(cls, config_path: Optional[str] = None):
|
||||
"""Load configuration settings from a JSON file.
|
||||
|
||||
|
||||
Reads configuration from the specified JSON file and updates the current
|
||||
configuration attributes. If the file doesn't exist or contains invalid
|
||||
JSON, logs an appropriate message and continues with default values.
|
||||
|
||||
|
||||
Args:
|
||||
config_path (str): Path to the JSON configuration file to load.
|
||||
Can be relative or absolute path.
|
||||
|
||||
Can be relative or absolute path. Defaults to
|
||||
:meth:`get_config_path`.
|
||||
|
||||
Raises:
|
||||
The method handles all exceptions internally and logs errors rather
|
||||
than propagating them, ensuring the system continues with defaults.
|
||||
|
||||
|
||||
Example:
|
||||
PromptManagerConfig.load_from_file('custom_config.json')
|
||||
PromptManagerConfig.load_from_file('/path/to/config.json')
|
||||
"""
|
||||
import json
|
||||
|
||||
|
||||
if config_path is None:
|
||||
config_path = cls.get_config_path()
|
||||
|
||||
if os.path.exists(config_path):
|
||||
try:
|
||||
with open(config_path, 'r') as f:
|
||||
with open(config_path, "r") as f:
|
||||
config = json.load(f)
|
||||
cls.update_config(config)
|
||||
config_logger.info(f"Loaded configuration from {config_path}")
|
||||
@@ -304,52 +617,50 @@ class PromptManagerConfig:
|
||||
config_logger.error(f"Error loading config from {config_path}: {e}")
|
||||
else:
|
||||
config_logger.info(f"Config file not found: {config_path}, using defaults")
|
||||
|
||||
|
||||
@classmethod
|
||||
def save_to_file(cls, config_path: str):
|
||||
def save_to_file(cls, config_path: Optional[str] = None):
|
||||
"""Save the current configuration to a JSON file.
|
||||
|
||||
|
||||
Serializes the complete configuration (including gallery settings) to
|
||||
a JSON file. Creates the directory structure if it doesn't exist.
|
||||
|
||||
|
||||
Args:
|
||||
config_path (str): Path where the JSON configuration file should be saved.
|
||||
Parent directories will be created if they don't exist.
|
||||
|
||||
Defaults to :meth:`get_config_path`.
|
||||
|
||||
Raises:
|
||||
The method handles all exceptions internally and logs errors rather
|
||||
than propagating them.
|
||||
|
||||
|
||||
Example:
|
||||
PromptManagerConfig.save_to_file('backup_config.json')
|
||||
PromptManagerConfig.save_to_file('/etc/comfyui/prompt_manager.json')
|
||||
"""
|
||||
import json
|
||||
|
||||
if config_path is None:
|
||||
config_path = cls.get_config_path()
|
||||
|
||||
try:
|
||||
config = cls.get_config()
|
||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
||||
|
||||
with open(config_path, 'w') as f:
|
||||
json.dump(config, f, indent=2)
|
||||
|
||||
write_private_json(config_path, cls.get_config())
|
||||
config_logger.info(f"Saved configuration to {config_path}")
|
||||
except Exception as e:
|
||||
config_logger.error(f"Error saving config to {config_path}: {e}")
|
||||
|
||||
|
||||
@classmethod
|
||||
def update_config(cls, new_config: Dict[str, Any]):
|
||||
"""Update configuration attributes from a dictionary.
|
||||
|
||||
|
||||
Takes a nested dictionary with configuration sections and updates
|
||||
the corresponding class attributes. Only updates attributes that
|
||||
are present in the input dictionary, leaving others unchanged.
|
||||
|
||||
|
||||
Args:
|
||||
new_config (Dict[str, Any]): Nested dictionary containing configuration
|
||||
updates. Should follow the same structure as returned by get_config().
|
||||
Valid top-level keys are: 'database', 'web_ui', 'performance', 'gallery'.
|
||||
|
||||
Valid top-level keys are: 'database', 'web_ui', 'performance',
|
||||
'gallery'.
|
||||
|
||||
Example:
|
||||
new_settings = {
|
||||
'database': {'default_path': 'custom.db'},
|
||||
@@ -357,39 +668,47 @@ class PromptManagerConfig:
|
||||
}
|
||||
PromptManagerConfig.update_config(new_settings)
|
||||
"""
|
||||
database = new_config.get('database', {})
|
||||
if 'default_path' in database:
|
||||
cls.DEFAULT_DB_PATH = database['default_path']
|
||||
if 'enable_duplicate_detection' in database:
|
||||
cls.ENABLE_DUPLICATE_DETECTION = database['enable_duplicate_detection']
|
||||
if 'enable_auto_save' in database:
|
||||
cls.ENABLE_AUTO_SAVE = database['enable_auto_save']
|
||||
|
||||
web_ui = new_config.get('web_ui', {})
|
||||
if 'result_timeout' in web_ui:
|
||||
cls.RESULT_TIMEOUT = web_ui['result_timeout']
|
||||
if 'show_test_button' in web_ui:
|
||||
cls.SHOW_TEST_BUTTON = web_ui['show_test_button']
|
||||
if 'webui_display_mode' in web_ui:
|
||||
cls.WEBUI_DISPLAY_MODE = web_ui['webui_display_mode']
|
||||
|
||||
performance = new_config.get('performance', {})
|
||||
if 'max_search_results' in performance:
|
||||
cls.MAX_SEARCH_RESULTS = performance['max_search_results']
|
||||
if 'enable_fuzzy_search' in performance:
|
||||
cls.ENABLE_FUZZY_SEARCH = performance['enable_fuzzy_search']
|
||||
if 'auto_backup_interval' in performance:
|
||||
cls.AUTO_BACKUP_INTERVAL = performance['auto_backup_interval']
|
||||
|
||||
database = new_config.get("database", {})
|
||||
if "default_path" in database:
|
||||
cls.DEFAULT_DB_PATH = database["default_path"]
|
||||
if "enable_duplicate_detection" in database:
|
||||
cls.ENABLE_DUPLICATE_DETECTION = database["enable_duplicate_detection"]
|
||||
if "enable_auto_save" in database:
|
||||
cls.ENABLE_AUTO_SAVE = database["enable_auto_save"]
|
||||
|
||||
web_ui = new_config.get("web_ui", {})
|
||||
if "result_timeout" in web_ui:
|
||||
cls.RESULT_TIMEOUT = web_ui["result_timeout"]
|
||||
if "show_test_button" in web_ui:
|
||||
cls.SHOW_TEST_BUTTON = web_ui["show_test_button"]
|
||||
if "webui_display_mode" in web_ui:
|
||||
cls.WEBUI_DISPLAY_MODE = web_ui["webui_display_mode"]
|
||||
if isinstance(web_ui.get("infinite_scroll"), bool):
|
||||
cls.INFINITE_SCROLL = web_ui["infinite_scroll"]
|
||||
|
||||
performance = new_config.get("performance", {})
|
||||
if "max_search_results" in performance:
|
||||
cls.MAX_SEARCH_RESULTS = performance["max_search_results"]
|
||||
if "enable_fuzzy_search" in performance:
|
||||
cls.ENABLE_FUZZY_SEARCH = performance["enable_fuzzy_search"]
|
||||
if "auto_backup_interval" in performance:
|
||||
cls.AUTO_BACKUP_INTERVAL = performance["auto_backup_interval"]
|
||||
if "worker_threads" in performance:
|
||||
cls.WORKER_THREADS = cls.clamp_worker_threads(
|
||||
performance["worker_threads"], cls.WORKER_THREADS
|
||||
)
|
||||
|
||||
# Update gallery config
|
||||
if 'gallery' in new_config:
|
||||
GalleryConfig.update_config(new_config['gallery'])
|
||||
if "gallery" in new_config:
|
||||
GalleryConfig.update_config(new_config["gallery"])
|
||||
|
||||
# Update integration config
|
||||
if "integrations" in new_config:
|
||||
IntegrationConfig.update_config(new_config["integrations"])
|
||||
|
||||
|
||||
# Load configuration on import
|
||||
try:
|
||||
config_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
config_file = os.path.join(config_dir, 'config.json')
|
||||
PromptManagerConfig.load_from_file(config_file)
|
||||
PromptManagerConfig.load_from_file()
|
||||
except Exception as e:
|
||||
config_logger.error(f"Error during config initialization: {e}")
|
||||
config_logger.error(f"Error during config initialization: {e}")
|
||||
|
||||
@@ -0,0 +1,706 @@
|
||||
"""Utilities for LoraManager integration.
|
||||
|
||||
Provides detection, metadata reading, and trigger word lookup for
|
||||
ComfyUI-Lora-Manager (https://github.com/willmiao/ComfyUI-Lora-Manager).
|
||||
|
||||
All functions are safe to call when LoraManager is not installed — they
|
||||
return empty results rather than raising.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
logger = get_logger("prompt_manager.lora_utils")
|
||||
|
||||
# ── LoraManager detection ────────────────────────────────────────────
|
||||
|
||||
|
||||
def find_comfyui_root() -> Optional[Path]:
|
||||
"""Walk upward from this file to find the ComfyUI root (contains main.py).
|
||||
|
||||
Tries both the resolved (real) path and the unresolved path to handle
|
||||
symlinked custom_nodes installations.
|
||||
"""
|
||||
start_paths = [Path(__file__).resolve().parent]
|
||||
|
||||
# If installed via symlink, the unresolved path leads through custom_nodes/
|
||||
raw_path = Path(__file__).parent
|
||||
if raw_path.resolve() != raw_path:
|
||||
start_paths.append(raw_path)
|
||||
|
||||
# Also try via folder_paths if available (ComfyUI runtime)
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
base = Path(folder_paths.base_path)
|
||||
if base.is_dir():
|
||||
return base
|
||||
except (ImportError, AttributeError):
|
||||
# folder_paths unavailable — not running inside ComfyUI runtime
|
||||
pass
|
||||
|
||||
for start in start_paths:
|
||||
current = start
|
||||
for _ in range(10):
|
||||
if (current / "main.py").exists() and (current / "custom_nodes").exists():
|
||||
return current
|
||||
parent = current.parent
|
||||
if parent == current:
|
||||
break
|
||||
current = parent
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def detect_lora_manager(custom_path: str = "") -> Optional[str]:
|
||||
"""Return the absolute path to ComfyUI-Lora-Manager if installed.
|
||||
|
||||
Args:
|
||||
custom_path: User-provided override path. Checked first.
|
||||
|
||||
Returns:
|
||||
Absolute path string, or None if not found.
|
||||
"""
|
||||
# 1. User override
|
||||
if custom_path:
|
||||
p = Path(custom_path)
|
||||
if p.is_dir() and _looks_like_lora_manager(p):
|
||||
return str(p.resolve())
|
||||
|
||||
# 2. Auto-detect via custom_nodes (case-insensitive scan)
|
||||
root = find_comfyui_root()
|
||||
if root:
|
||||
custom_nodes = root / "custom_nodes"
|
||||
if custom_nodes.is_dir():
|
||||
for entry in custom_nodes.iterdir():
|
||||
if (
|
||||
entry.is_dir()
|
||||
and "lora" in entry.name.lower()
|
||||
and "manager" in entry.name.lower()
|
||||
and _looks_like_lora_manager(entry)
|
||||
):
|
||||
return str(entry.resolve())
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _looks_like_lora_manager(path: Path) -> bool:
|
||||
"""Heuristic: does this directory look like a LoraManager install?"""
|
||||
# Must have __init__.py (ComfyUI extension) or README.md
|
||||
has_init = (path / "__init__.py").exists()
|
||||
if not has_init:
|
||||
return False
|
||||
# Check for characteristic structure: py/ dir, or any .metadata.json nearby
|
||||
return (path / "py").is_dir() or (path / "lora_manager").is_dir()
|
||||
|
||||
|
||||
def _public_path_anchors() -> List[str]:
|
||||
"""Directories the API's public (relative) paths are rendered against."""
|
||||
try:
|
||||
from .config import GalleryConfig
|
||||
except ImportError:
|
||||
try:
|
||||
from config import GalleryConfig
|
||||
except ImportError:
|
||||
return []
|
||||
try:
|
||||
return list(GalleryConfig.path_anchors())
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def custom_nodes_directories() -> List[Path]:
|
||||
"""ComfyUI ``custom_nodes`` directories a LoraManager install may live in.
|
||||
|
||||
Inside ComfyUI these come from ``folder_paths`` (``base_path`` plus any
|
||||
registered ``custom_nodes`` folders); outside it, the parent of this
|
||||
package, which is ``custom_nodes`` for a normal install.
|
||||
"""
|
||||
found: List[Path] = []
|
||||
|
||||
def add(candidate) -> None:
|
||||
try:
|
||||
resolved = Path(candidate).resolve()
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
return
|
||||
if resolved not in found:
|
||||
found.append(resolved)
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
folder_paths = None
|
||||
if folder_paths is not None:
|
||||
base = getattr(folder_paths, "base_path", None)
|
||||
if isinstance(base, str) and base:
|
||||
add(Path(base) / "custom_nodes")
|
||||
getter = getattr(folder_paths, "get_folder_paths", None)
|
||||
if callable(getter):
|
||||
try:
|
||||
registered = getter("custom_nodes") or []
|
||||
except Exception:
|
||||
registered = []
|
||||
for entry in registered:
|
||||
add(entry)
|
||||
if not found:
|
||||
add(Path(__file__).resolve().parent.parent.parent)
|
||||
return found
|
||||
|
||||
|
||||
def resolve_lora_manager_path(path) -> Optional[str]:
|
||||
"""Canonical form of ``path`` when it may be used as the LoraManager install.
|
||||
|
||||
Accepted: an existing directory that is a direct child of one of
|
||||
:func:`custom_nodes_directories` and whose name contains ``lora``
|
||||
(case-insensitive). A relative path is tried against the ComfyUI root
|
||||
and each custom_nodes directory. Anything else, including a symlink
|
||||
that escapes custom_nodes, yields None.
|
||||
"""
|
||||
if not isinstance(path, str) or not path.strip():
|
||||
return None
|
||||
path = path.strip()
|
||||
roots = []
|
||||
for root in custom_nodes_directories():
|
||||
try:
|
||||
roots.append(Path(os.path.normcase(os.path.realpath(str(root)))))
|
||||
except (OSError, ValueError):
|
||||
continue
|
||||
|
||||
candidates = [path]
|
||||
if not os.path.isabs(path):
|
||||
# Public (API) paths are relative to these anchors; the ComfyUI root
|
||||
# and the custom_nodes directories cover hand-typed relative paths.
|
||||
bases = list(_public_path_anchors())
|
||||
comfy_root = find_comfyui_root()
|
||||
if comfy_root:
|
||||
bases.append(comfy_root)
|
||||
bases.extend(root.parent for root in roots)
|
||||
bases.extend(roots)
|
||||
candidates = [os.path.join(str(base), path) for base in bases]
|
||||
|
||||
for candidate in candidates:
|
||||
try:
|
||||
canonical = Path(os.path.normcase(os.path.realpath(candidate)))
|
||||
is_dir = canonical.is_dir()
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
continue
|
||||
if not is_dir or "lora" not in canonical.name.lower():
|
||||
continue
|
||||
if canonical.parent in roots:
|
||||
return str(canonical)
|
||||
return None
|
||||
|
||||
|
||||
# ── Metadata reading ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def find_lora_directories(lora_manager_path: str) -> List[str]:
|
||||
"""Find directories that contain LoRA models (with .metadata.json files).
|
||||
|
||||
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()))
|
||||
|
||||
# Also try folder_paths at runtime (catches all configured paths)
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
for p in folder_paths.get_folder_paths("loras"):
|
||||
pp = Path(p)
|
||||
if pp.is_dir():
|
||||
dirs.add(str(pp.resolve()))
|
||||
except (ImportError, AttributeError):
|
||||
# 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 _get_extra_lora_paths(comfyui_root: Path) -> List[Path]:
|
||||
"""Parse extra_model_paths.yaml for additional LoRA directories."""
|
||||
results = []
|
||||
for name in ("extra_model_paths.yaml", "extra_model_paths.yml"):
|
||||
config_file = comfyui_root / name
|
||||
if not config_file.exists():
|
||||
continue
|
||||
try:
|
||||
import yaml
|
||||
|
||||
config = yaml.safe_load(config_file.read_text())
|
||||
if not isinstance(config, dict):
|
||||
continue
|
||||
for section in config.values():
|
||||
if not isinstance(section, dict):
|
||||
continue
|
||||
base = Path(section.get("base_path", ""))
|
||||
loras_val = section.get("loras", "")
|
||||
if not loras_val:
|
||||
continue
|
||||
for line in str(loras_val).strip().splitlines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
p = Path(line)
|
||||
if not p.is_absolute():
|
||||
p = base / line
|
||||
if p.is_dir():
|
||||
results.append(p)
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to parse {config_file}: {e}")
|
||||
return results
|
||||
|
||||
|
||||
def read_lora_metadata(metadata_path: Path) -> Optional[Dict]:
|
||||
"""Read and parse a single .metadata.json file.
|
||||
|
||||
Returns:
|
||||
Parsed dict, or None on failure.
|
||||
"""
|
||||
try:
|
||||
with open(metadata_path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
logger.debug(f"Failed to read {metadata_path}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _get_civitai(metadata: Dict) -> Dict:
|
||||
"""Safely get the civitai dict, handling None values."""
|
||||
return metadata.get("civitai") or {}
|
||||
|
||||
|
||||
def get_trigger_words_from_metadata(metadata: Dict) -> List[str]:
|
||||
"""Extract trigger words from a parsed LoraManager metadata dict."""
|
||||
civitai = _get_civitai(metadata)
|
||||
words = civitai.get("trainedWords", [])
|
||||
if isinstance(words, list):
|
||||
return [w.strip() for w in words if isinstance(w, str) and w.strip()]
|
||||
return []
|
||||
|
||||
|
||||
def get_example_prompt_from_metadata(metadata: Dict) -> Optional[str]:
|
||||
"""Extract an example prompt from civitai image metadata.
|
||||
|
||||
Looks at civitai.images[].meta.prompt for the first available example.
|
||||
"""
|
||||
civitai = _get_civitai(metadata)
|
||||
images = civitai.get("images", []) or []
|
||||
for img in images:
|
||||
if not isinstance(img, dict):
|
||||
continue
|
||||
meta = img.get("meta")
|
||||
if isinstance(meta, dict):
|
||||
prompt = meta.get("prompt", "")
|
||||
if isinstance(prompt, str) and prompt.strip():
|
||||
return prompt.strip()
|
||||
return None
|
||||
|
||||
|
||||
def get_civitai_image_urls(metadata: Dict) -> List[str]:
|
||||
"""Extract all civitai example image URLs from metadata."""
|
||||
civitai = _get_civitai(metadata)
|
||||
urls = []
|
||||
for img in civitai.get("images", []) or []:
|
||||
if not isinstance(img, dict):
|
||||
continue
|
||||
url = img.get("url", "")
|
||||
if isinstance(url, str) and url.strip():
|
||||
urls.append(url.strip())
|
||||
return urls
|
||||
|
||||
|
||||
def get_model_name_from_metadata(metadata: Dict) -> str:
|
||||
"""Extract the model display name from metadata."""
|
||||
name = metadata.get("model_name", "")
|
||||
if not name:
|
||||
civitai = _get_civitai(metadata)
|
||||
model = civitai.get("model") or {}
|
||||
name = model.get("name", "")
|
||||
if not name:
|
||||
name = metadata.get("file_name", "unknown")
|
||||
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
|
||||
|
||||
# Only these hosts may ever receive the CivitAI API key, and only over HTTPS.
|
||||
CIVITAI_HOSTS = frozenset(
|
||||
{"civitai.com", "www.civitai.com", "api.civitai.com", "image.civitai.com"}
|
||||
)
|
||||
MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024
|
||||
_DOWNLOAD_CHUNK_BYTES = 64 * 1024
|
||||
|
||||
|
||||
class _RefuseRedirects(urllib.request.HTTPRedirectHandler):
|
||||
"""Never follow a redirect: urllib would re-send the Bearer header to it."""
|
||||
|
||||
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
||||
logger.warning(f"Refusing redirect from {req.full_url} to {newurl}")
|
||||
return None
|
||||
|
||||
|
||||
def _build_opener(*handlers) -> urllib.request.OpenerDirector:
|
||||
"""Build the outbound opener; extra handlers let tests serve canned responses."""
|
||||
return urllib.request.build_opener(_RefuseRedirects, *handlers)
|
||||
|
||||
|
||||
_opener = _build_opener()
|
||||
|
||||
|
||||
def _open_url(req: urllib.request.Request, timeout: float = 10):
|
||||
"""Open ``req`` through the redirect-refusing opener."""
|
||||
return _opener.open(req, timeout=timeout)
|
||||
|
||||
|
||||
def is_civitai_url(url: str) -> bool:
|
||||
"""True when ``url`` is an HTTPS URL on an allow-listed CivitAI host."""
|
||||
try:
|
||||
parts = urlsplit(url)
|
||||
port = parts.port
|
||||
except ValueError:
|
||||
return False
|
||||
hostname = (parts.hostname or "").lower()
|
||||
return (
|
||||
parts.scheme == "https"
|
||||
and hostname in CIVITAI_HOSTS
|
||||
and parts.username is None
|
||||
and parts.password is None
|
||||
and port in (None, 443)
|
||||
)
|
||||
|
||||
|
||||
def _read_capped(resp, cap: int) -> Optional[bytes]:
|
||||
"""Read a response body, refusing anything larger than ``cap`` bytes."""
|
||||
declared = resp.headers.get("Content-Length")
|
||||
if declared is not None:
|
||||
try:
|
||||
if int(declared) > cap:
|
||||
return None
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
chunks = []
|
||||
total = 0
|
||||
while True:
|
||||
chunk = resp.read(_DOWNLOAD_CHUNK_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
total += len(chunk)
|
||||
if total > cap:
|
||||
return None
|
||||
chunks.append(chunk)
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
|
||||
"""Download a single CivitAI image, resize to thumbnail, save as JPEG.
|
||||
|
||||
Refuses any URL that is not HTTPS on an allow-listed CivitAI host so the
|
||||
API key is never sent elsewhere, and refuses bodies larger than
|
||||
MAX_CIVITAI_DOWNLOAD_BYTES.
|
||||
"""
|
||||
if not is_civitai_url(url):
|
||||
logger.warning(f"Refusing non-CivitAI download URL: {url}")
|
||||
return None
|
||||
try:
|
||||
headers = {"User-Agent": "ComfyUI-PromptManager/1.0"}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
req = urllib.request.Request(url, headers=headers)
|
||||
with _open_url(req, timeout=10) as resp:
|
||||
raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
|
||||
if raw is None:
|
||||
logger.warning(f"Refusing oversized download: {url}")
|
||||
return None
|
||||
|
||||
# Resize to thumbnail to save disk space
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image
|
||||
|
||||
img = Image.open(BytesIO(raw))
|
||||
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
|
||||
img = img.convert("RGB")
|
||||
img.save(str(local_path), "JPEG", quality=85)
|
||||
|
||||
return str(local_path.resolve())
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to download {url}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def download_civitai_images(
|
||||
metadata: Dict, metadata_path: Path, cache_dir: Path, api_key: str = ""
|
||||
) -> List[str]:
|
||||
"""Download civitai example images to a local cache directory.
|
||||
|
||||
Uses thumbnail URLs (512px) instead of full-size originals, and
|
||||
downloads in parallel (up to 8 concurrent) for speed.
|
||||
|
||||
Args:
|
||||
api_key: CivitAI API key for authenticated downloads (NSFW content).
|
||||
|
||||
Returns:
|
||||
List of absolute paths to downloaded image files.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
civitai = _get_civitai(metadata)
|
||||
images = civitai.get("images", []) or []
|
||||
if not images:
|
||||
return []
|
||||
|
||||
file_name = metadata.get("file_name", "")
|
||||
if not file_name:
|
||||
file_name = metadata_path.name.replace(".metadata.json", "")
|
||||
lora_stem = Path(file_name).stem
|
||||
|
||||
lora_cache = cache_dir / lora_stem
|
||||
lora_cache.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Build download tasks
|
||||
cached = []
|
||||
tasks = [] # (url, local_path)
|
||||
for img in images:
|
||||
if not isinstance(img, dict):
|
||||
continue
|
||||
url = img.get("url", "")
|
||||
if not isinstance(url, str) or not is_civitai_url(url):
|
||||
continue
|
||||
|
||||
url_hash = hashlib.md5(url.encode()).hexdigest()[:12]
|
||||
local_path = lora_cache / f"{url_hash}.jpg"
|
||||
|
||||
if local_path.exists():
|
||||
cached.append(str(local_path.resolve()))
|
||||
else:
|
||||
tasks.append((url, local_path))
|
||||
|
||||
if not tasks:
|
||||
return cached
|
||||
|
||||
# Download in parallel
|
||||
downloaded = []
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
futures = [
|
||||
pool.submit(_download_one, url, path, api_key) for url, path in tasks
|
||||
]
|
||||
for fut in futures:
|
||||
result = fut.result()
|
||||
if result:
|
||||
downloaded.append(result)
|
||||
|
||||
return cached + downloaded
|
||||
|
||||
|
||||
def get_lora_image_cache_dir() -> Path:
|
||||
"""Get the directory used to cache downloaded LoRA example images."""
|
||||
# Store in the extension's own directory
|
||||
ext_root = Path(__file__).resolve().parent.parent
|
||||
cache = ext_root / "data" / "lora_images"
|
||||
cache.mkdir(parents=True, exist_ok=True)
|
||||
return cache
|
||||
|
||||
|
||||
def get_example_images_dir(lora_manager_path: str) -> Optional[str]:
|
||||
"""Find the LoraManager example_images directory."""
|
||||
lm_path = Path(lora_manager_path)
|
||||
|
||||
# Direct subdirectory
|
||||
candidate = lm_path / "example_images"
|
||||
if candidate.is_dir():
|
||||
return str(candidate.resolve())
|
||||
|
||||
# Search one level in user data dirs
|
||||
for child in lm_path.iterdir():
|
||||
if child.is_dir():
|
||||
sub = child / "example_images"
|
||||
if sub.is_dir():
|
||||
return str(sub.resolve())
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ── Trigger word cache & injection ───────────────────────────────────
|
||||
|
||||
_LORA_PATTERN = re.compile(r"<lora:([^:>]+):[^>]+>", re.IGNORECASE)
|
||||
|
||||
|
||||
class TriggerWordCache:
|
||||
"""Thread-safe cache mapping LoRA names to their trigger words.
|
||||
|
||||
Built lazily on first access, refreshable on demand.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._cache: Dict[str, List[str]] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._loaded = False
|
||||
|
||||
def load(self, lora_manager_path: str) -> int:
|
||||
"""Scan LoRA metadata files and build the trigger word mapping.
|
||||
|
||||
Returns:
|
||||
Number of LoRAs with trigger words found.
|
||||
"""
|
||||
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
|
||||
|
||||
logger.info(
|
||||
f"Trigger word cache loaded: {len(new_cache)} LoRAs with trigger words"
|
||||
)
|
||||
return len(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(), [])
|
||||
|
||||
@property
|
||||
def is_loaded(self) -> bool:
|
||||
with self._lock:
|
||||
return self._loaded
|
||||
|
||||
def clear(self):
|
||||
with self._lock:
|
||||
self._cache.clear()
|
||||
self._loaded = False
|
||||
|
||||
|
||||
# Module-level singleton
|
||||
_trigger_cache = TriggerWordCache()
|
||||
|
||||
|
||||
def get_trigger_cache() -> TriggerWordCache:
|
||||
return _trigger_cache
|
||||
|
||||
|
||||
def inject_trigger_words(text: str, cache: TriggerWordCache) -> Tuple[str, List[str]]:
|
||||
"""Scan text for <lora:NAME:WEIGHT> tags and append trigger words.
|
||||
|
||||
Args:
|
||||
text: The prompt text potentially containing lora tags.
|
||||
cache: Populated TriggerWordCache instance.
|
||||
|
||||
Returns:
|
||||
Tuple of (modified_text, list_of_injected_words).
|
||||
If no trigger words found, returns the original text unchanged.
|
||||
"""
|
||||
if not cache.is_loaded:
|
||||
return text, []
|
||||
|
||||
matches = _LORA_PATTERN.findall(text)
|
||||
if not matches:
|
||||
return text, []
|
||||
|
||||
all_words = []
|
||||
for lora_name in matches:
|
||||
words = cache.get_trigger_words(lora_name)
|
||||
for w in words:
|
||||
if w.lower() not in text.lower() and w not in all_words:
|
||||
all_words.append(w)
|
||||
|
||||
if not all_words:
|
||||
return text, []
|
||||
|
||||
injected = ", ".join(all_words)
|
||||
return f"{text}, {injected}", all_words
|
||||
@@ -1,9 +1,14 @@
|
||||
[project]
|
||||
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.0.21"
|
||||
version = "3.3.1"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["# Core dependencies for PromptManager", "# Note: Most dependencies are already included with ComfyUI", "# Already included with Python standard library:", "# - sqlite3", "# - hashlib", "# - json", "# - datetime", "# - os", "# - typing", "# - threading", "# - uuid", "# Required for gallery functionality:", "watchdog>=2.1.0 # For file system monitoring", "Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)", "# Optional dependencies for enhanced search functionality:", "# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)", "# sqlalchemy>=1.4.0 # For advanced ORM features (optional)", "# Development dependencies (optional):", "# pytest>=6.0.0 # For running tests", "# black>=22.0.0 # For code formatting", "# flake8>=4.0.0 # For linting", "# mypy>=0.910 # For type checking"]
|
||||
requires-python = ">=3.10"
|
||||
dependencies = ["watchdog>=2.1.0", "Pillow>=8.0.0"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=84.0.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/ComfyAssets/ComfyUI_PromptManager"
|
||||
@@ -13,3 +18,10 @@ Repository = "https://github.com/ComfyAssets/ComfyUI_PromptManager"
|
||||
PublisherId = "kiko9"
|
||||
DisplayName = "ComfyUI_PromptManager"
|
||||
Icon = "https://avatars.githubusercontent.com/u/213204677?s=200"
|
||||
|
||||
[tool.black]
|
||||
line-length = 88
|
||||
target-version = ["py310"]
|
||||
|
||||
[tool.bandit]
|
||||
exclude_dirs = ["tests", "reference"]
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# Runtime dependencies
|
||||
watchdog>=6.0.0
|
||||
Pillow>=12.3.0
|
||||
aiohttp>=3.14.5
|
||||
|
||||
# Tests and coverage
|
||||
coverage>=7.16.2
|
||||
numpy>=2.2.6 # WD14 engine tests
|
||||
|
||||
# Code quality
|
||||
black>=26.10.1
|
||||
flake8>=7.4.1
|
||||
|
||||
# Security scanning
|
||||
bandit[toml]>=1.9.4
|
||||
|
||||
# Pre-commit hooks
|
||||
pre-commit>=4.6.2
|
||||
@@ -1,26 +1,4 @@
|
||||
# Core dependencies for PromptManager
|
||||
# Note: Most dependencies are already included with ComfyUI
|
||||
|
||||
# Already included with Python standard library:
|
||||
# - sqlite3
|
||||
# - hashlib
|
||||
# - json
|
||||
# - datetime
|
||||
# - os
|
||||
# - typing
|
||||
# - threading
|
||||
# - uuid
|
||||
|
||||
# Required for gallery functionality:
|
||||
watchdog>=2.1.0 # For file system monitoring
|
||||
Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)
|
||||
|
||||
# Optional dependencies for enhanced search functionality:
|
||||
# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)
|
||||
# sqlalchemy>=1.4.0 # For advanced ORM features (optional)
|
||||
|
||||
# Development dependencies (optional):
|
||||
# pytest>=6.0.0 # For running tests
|
||||
# black>=22.0.0 # For code formatting
|
||||
# flake8>=4.0.0 # For linting
|
||||
# mypy>=0.910 # For type checking
|
||||
# Runtime dependencies for PromptManager.
|
||||
# aiohttp, numpy and torch come from ComfyUI itself and are not listed here.
|
||||
watchdog>=2.1.0 # Output directory monitoring
|
||||
Pillow>=8.0.0 # Image metadata extraction and thumbnails
|
||||
|
||||
@@ -1,63 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ComfyUI_PromptManager Gallery System Restart Script.
|
||||
|
||||
This script reinitializes and restarts the automatic image gallery monitoring system
|
||||
for ComfyUI_PromptManager. Use this script after code modifications or when the
|
||||
gallery system needs to be reloaded.
|
||||
|
||||
The script performs the following operations:
|
||||
1. Initializes the database connection
|
||||
2. Creates a PromptTracker instance for linking prompts to images
|
||||
3. Sets up the ImageMonitor for real-time file system monitoring
|
||||
4. Starts monitoring ComfyUI output directories
|
||||
5. Reports system status and readiness
|
||||
|
||||
Usage:
|
||||
python restart_gallery.py
|
||||
|
||||
This is particularly useful during development when you need to restart the
|
||||
gallery monitoring system without restarting the entire ComfyUI application.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add current directory to Python path
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.insert(0, current_dir)
|
||||
|
||||
try:
|
||||
# Import and initialize the gallery system
|
||||
from database.operations import PromptDatabase
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
from utils.image_monitor import ImageMonitor
|
||||
|
||||
print("[RESTART] Restarting PromptManager Gallery System...")
|
||||
|
||||
# Initialize components
|
||||
db = PromptDatabase()
|
||||
tracker = PromptTracker(db)
|
||||
monitor = ImageMonitor(db, tracker)
|
||||
|
||||
print("[SUCCESS] Components initialized successfully")
|
||||
|
||||
# Test image monitoring directories
|
||||
status = monitor.get_status()
|
||||
print(f"[INFO] Monitoring status: {status}")
|
||||
|
||||
# Start monitoring
|
||||
monitor.start_monitoring()
|
||||
print("[SUCCESS] Image monitoring restarted")
|
||||
|
||||
print("\n[READY] Gallery system restart completed!")
|
||||
print("Generate some images now to test the automatic linking.")
|
||||
|
||||
except ImportError as e:
|
||||
print(f"[ERROR] Import error: {e}")
|
||||
print("Make sure you're running this from the PromptManager directory")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
@@ -3,4 +3,4 @@ Test suite for ComfyUI_PromptManager.
|
||||
|
||||
This package contains unit tests and integration tests for the PromptManager
|
||||
custom nodes, database operations, utility functions, and web interface components.
|
||||
"""
|
||||
"""
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Browser workflow fixtures; never imported by production code."""
|
||||
@@ -0,0 +1,150 @@
|
||||
import { readFileSync } from 'node:fs';
|
||||
import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
|
||||
test.beforeEach(async ({page, browserName}) => {
|
||||
await startCoverage(page, browserName);
|
||||
page.classicErrors = [];
|
||||
page.on('pageerror', error => page.classicErrors.push(error.message));
|
||||
await page.route('**/scripts/app.js', route => route.fulfill({contentType:'text/javascript', body:'export const app = window.fixtureApp;'}));
|
||||
await page.route('**/scripts/api.js', route => route.fulfill({contentType:'text/javascript', body:'export const api = {fileURL: path => path, queuePrompt: async (...args) => args};'}));
|
||||
await page.route('**/extensions/promptmanager/prompt_manager.js', route => route.fulfill({contentType:'text/javascript', body:readFileSync(new URL('../../web/comfy/prompt_manager.js', import.meta.url),'utf8')}));
|
||||
await page.goto('/prompt_manager/test');
|
||||
await page.evaluate(() => {
|
||||
document.body.replaceChildren();
|
||||
document.body.style.setProperty('--bg-color','#222');
|
||||
window.fixtureApp={extensions:[],registerExtension(extension){this.extensions.push(extension);},ui:{settings:new EventTarget()},canvas:{},graph:{setDirtyCanvas(){}}};
|
||||
});
|
||||
await page.evaluate(async () => {
|
||||
await import('/extensions/promptmanager/prompt_manager.js');
|
||||
class Node {
|
||||
constructor(){
|
||||
this.widgets=['text','category','tags','search_text'].map(name=>({name,value:'',inputEl:document.createElement('textarea')}));
|
||||
this.widgets_values=[]; this.inputs=[];
|
||||
}
|
||||
addDOMWidget(name,type,element){document.body.append(element);return {element};}
|
||||
setSize(){} setDirtyCanvas(){} serialize(){return {properties:{...this.properties}};}
|
||||
}
|
||||
await fixtureApp.extensions[0].beforeRegisterNodeDef(Node,{name:'PromptManager'},fixtureApp);
|
||||
window.node=new Node();node.onNodeCreated();
|
||||
});
|
||||
await expect(page.locator('.pm-classic')).toBeVisible();
|
||||
});
|
||||
|
||||
test.afterEach(async ({page,browserName}) => {
|
||||
await saveCoverage(page,browserName);
|
||||
expect(page.classicErrors).toEqual([]);
|
||||
});
|
||||
|
||||
test('Classic keeps immediate selection, search query and metadata serialization',async({page})=>{
|
||||
await page.route('**/prompt_manager/search?**', route=>route.fulfill({json:{success:true,results:[{text:'A new portrait',category:'Portrait',tags:['film grain'],rating:5}]}}));
|
||||
await page.evaluate(()=>{node.widgets.find(w=>w.name==='search_text').value='tag:"film grain" portrait';});
|
||||
const response=page.waitForRequest('**/prompt_manager/search?**');
|
||||
await page.getByRole('button',{name:'🔍 Search',exact:true}).click();
|
||||
expect(new URL((await response).url()).searchParams.get('tags')).toBe('film grain');
|
||||
const result=page.getByRole('button',{name:/A new portrait/});
|
||||
await result.focus();await result.press('Enter');
|
||||
expect(await page.evaluate(()=>node.serialize().widgets_values.slice(0,3))).toEqual(['A new portrait','Portrait','film grain']);
|
||||
await expect(result).toHaveAttribute('aria-pressed','true');
|
||||
const released=await page.evaluate(()=>{
|
||||
const previous=node.resultsSection.querySelector('button');
|
||||
node.displayResults([{text:'Replacement',tags:[]}]);
|
||||
previous.click();
|
||||
return node.widgets.find(widget=>widget.name==='text').value;
|
||||
});
|
||||
expect(released).toBe('A new portrait');
|
||||
await expect(page.getByRole('button',{name:'🌐 Web UI',exact:true})).toBeVisible();
|
||||
});
|
||||
|
||||
test('Classic detaches native text synchronization on removal',async({page})=>{
|
||||
await page.waitForTimeout(60);
|
||||
const values=await page.evaluate(()=>{
|
||||
const widget=node.widgets.find(item=>item.name==='text');
|
||||
widget.inputEl.value='Before removal';widget.inputEl.dispatchEvent(new Event('input'));
|
||||
const before=widget.value;node.onRemoved();
|
||||
widget.inputEl.value='After removal';widget.inputEl.dispatchEvent(new Event('input'));
|
||||
return [before,widget.value];
|
||||
});
|
||||
expect(values).toEqual(['Before removal','Before removal']);
|
||||
});
|
||||
|
||||
test('Classic follows host themes, keeps restore behavior and disposes resources',async({page})=>{
|
||||
const panel=page.locator('.pm-classic');
|
||||
await expect(panel).toHaveAttribute('data-ca-theme','dark');
|
||||
await page.evaluate(()=>document.body.style.setProperty('--bg-color','#fff'));
|
||||
await expect(panel).toHaveAttribute('data-ca-theme','light');
|
||||
await page.evaluate(()=>{node.properties.resultTimeout=0;node.displayResults([{text:'Example',tags:[]}]);node.hideResults();});
|
||||
await page.getByRole('button',{name:/Results auto-hidden/}).click();
|
||||
await expect(page.getByRole('button',{name:/Example/})).toBeVisible();
|
||||
const state=await page.evaluate(()=>{
|
||||
node.showNotification('Failure','error');node.onRemoved();
|
||||
const disposed=node._classicScope.signal.aborted;
|
||||
node.showNotification('Late response','info');
|
||||
return {disposed,notifications:document.querySelectorAll('.pm-classic-notification').length,style:node._styleSheet.isConnected};
|
||||
});
|
||||
expect(state).toEqual({disposed:true,notifications:0,style:false});
|
||||
await page.evaluate(()=>document.body.style.setProperty('--bg-color','#222'));
|
||||
await expect(panel).toHaveAttribute('data-ca-theme','light');
|
||||
});
|
||||
|
||||
test('Classic handles empty queries, empty results, failed requests and late responses',async({page})=>{
|
||||
await page.getByRole('button',{name:'🔍 Search',exact:true}).click();
|
||||
await expect(page.locator('.pm-classic-results')).toContainText('Enter search text');
|
||||
await page.route('**/prompt_manager/recent?**',route=>route.fulfill({json:{success:true,results:[]}}));
|
||||
await page.getByRole('button',{name:'📋 Recent',exact:true}).click();
|
||||
await expect(page.locator('.pm-classic-results')).toContainText('No prompts found');
|
||||
await page.evaluate(async()=>{
|
||||
node.callNodeMethod=async()=>{throw new Error('<img src=x onerror=alert(1)>');};
|
||||
await node.loadRecentPrompts();
|
||||
});
|
||||
await expect(page.locator('.pm-classic-results [role="alert"]')).toContainText('<img src=x');
|
||||
await expect(page.locator('.pm-classic-results img')).toHaveCount(0);
|
||||
const late=await page.evaluate(async()=>{
|
||||
let finish; node.callNodeMethod=()=>new Promise(resolve=>{finish=resolve;});
|
||||
const loading=node.loadRecentPrompts();node.onRemoved();finish({results:[{text:'Late'}]});await loading;
|
||||
return {notifications:node._classicNotifications.size,late:node.resultsSection.textContent.includes('Late')};
|
||||
});
|
||||
expect(late).toEqual({notifications:0,late:false});
|
||||
});
|
||||
|
||||
test('shipped shared factories keep keyboard behavior and theme isolation across mounts',async({page})=>{
|
||||
await page.evaluate(async()=>{
|
||||
const ui=await import('/prompt_manager/js/shared/toolkit/core.js');
|
||||
const scope=ui.createScope();const panel=document.createElement('div');panel.className='ca-tools';panel.id='shared-contract';
|
||||
window.shared={ui,scope,panel,events:[]};
|
||||
document.body.append(panel);
|
||||
ui.mountTheme(panel,null,scope);
|
||||
const save=ui.createButton('Save contract','ca-primary');scope.listen(save,'click',()=>shared.events.push('save'));
|
||||
panel.append(save,ui.createSlider('Size',{id:'contract-size',min:0,max:100,step:1,value:50,unit:'px'},scope,value=>shared.events.push(value)),ui.createSwitch('Enabled','','contract-enabled'),ui.createField('Order',ui.createSelect([['name','Name'],['date','Date']],''),'contract-order'),ui.createSegments('View',[['grid','Grid'],['list','List']],'contract-view',scope,value=>shared.events.push(value)));
|
||||
});
|
||||
await page.getByRole('button',{name:'Save contract',exact:true}).press('Enter');
|
||||
await page.getByRole('slider',{name:'Size'}).press('ArrowRight');
|
||||
await expect(page.locator('#shared-contract output')).toHaveText('51 px');
|
||||
await page.getByRole('switch',{name:'Enabled'}).uncheck();
|
||||
await page.getByLabel('Order').selectOption('date');
|
||||
await page.getByRole('radio',{name:'List',exact:true}).check();
|
||||
expect(await page.evaluate(()=>shared.events)).toEqual(['save',51,'list']);
|
||||
await page.evaluate(()=>{document.body.style.removeProperty('--bg-color');fixtureApp.ui.settings.dispatchEvent(new Event('Comfy.ColorPalette.change'));});
|
||||
await page.emulateMedia({colorScheme:'light'});
|
||||
await expect(page.locator('#shared-contract')).toHaveAttribute('data-ca-theme','light');
|
||||
await page.emulateMedia({colorScheme:'dark'});
|
||||
await expect(page.locator('#shared-contract')).toHaveAttribute('data-ca-theme','dark');
|
||||
await page.evaluate(()=>{shared.scope.dispose();document.body.style.setProperty('--bg-color','#fff');});
|
||||
await expect(page.locator('#shared-contract')).toHaveAttribute('data-ca-theme','dark');
|
||||
await expect(page.locator('.pm-classic')).toHaveAttribute('data-ca-theme','light');
|
||||
});
|
||||
|
||||
test('Classic test button reports success and failure and releases expired notifications',async({page})=>{
|
||||
await page.clock.install();
|
||||
await page.evaluate(()=>{node.properties.showTestButton=true;node.createSearchUI(node._classicContainer);});
|
||||
await page.getByRole('button',{name:'🔧 Test API',exact:true}).click();
|
||||
await expect(page.locator('.pm-classic-results')).toContainText('API Connection Successful');
|
||||
await page.route('**/prompt_manager/test',route=>route.fulfill({status:503,body:'unavailable'}));
|
||||
await page.getByRole('button',{name:'🔧 Test API',exact:true}).click();
|
||||
await expect(page.locator('.pm-classic-results [role="alert"]')).toContainText('503');
|
||||
await page.route('**/prompt_manager/test',route=>route.fulfill({json:{success:false,message:'Not ready'}}));
|
||||
await page.getByRole('button',{name:'🔧 Test API',exact:true}).click();
|
||||
await expect(page.locator('.pm-classic-results [role="alert"]')).toContainText('Not ready');
|
||||
await page.clock.fastForward(3001);
|
||||
await expect(page.locator('.pm-classic-notification')).toHaveCount(0);
|
||||
});
|
||||
@@ -0,0 +1,411 @@
|
||||
import {readFileSync} from 'node:fs';
|
||||
import {test,expect} from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import {startCoverage,saveCoverage} from '../../tools/vendor/browser-coverage.mjs';
|
||||
|
||||
test.beforeEach(async({page,browserName})=>{
|
||||
await page.addInitScript(()=>Object.defineProperty(Crypto.prototype,'randomUUID',{value:undefined,configurable:true}));
|
||||
await startCoverage(page,browserName);
|
||||
page.focusedErrors=[];page.on('pageerror',error=>page.focusedErrors.push(error.message));
|
||||
await page.route('**/scripts/app.js',route=>route.fulfill({contentType:'text/javascript',body:'export const app=window.fixtureApp;'}));
|
||||
await page.route('**/scripts/api.js',route=>route.fulfill({contentType:'text/javascript',body:'export const api={fileURL:path=>path,queuePrompt:async(...args)=>args};'}));
|
||||
await page.route('**/extensions/promptmanager/prompt_manager.js',route=>route.fulfill({contentType:'text/javascript',body:readFileSync(new URL('../../web/comfy/prompt_manager.js',import.meta.url),'utf8')}));
|
||||
await page.goto('/prompt_manager/test');
|
||||
await page.evaluate(async()=>{
|
||||
document.body.replaceChildren();document.body.style.setProperty('--bg-color','#222');
|
||||
window.fixtureApp={registerExtension(extension){this.extension=extension;},ui:{settings:new EventTarget()},graph:{setDirtyCanvas(){}},canvas:{}};
|
||||
await import('/extensions/promptmanager/prompt_manager.js');
|
||||
class Node {
|
||||
constructor(){this.widgets=['text','category','tags','search_text'].map(name=>({name,type:'text',value:'',options:{},inputEl:document.createElement('textarea')}));this.inputs=['clip','prepend_text','append_text'].map(name=>({name,link:null}));this.properties={};this.size=[550,700];}
|
||||
addWidget(type,name,value,callback,options){const widget={type,name,value,callback,options};this.widgets.push(widget);return widget;}
|
||||
addDOMWidget(name,type,element){document.body.append(element);return {element};}
|
||||
setSize(size){this.size=size;}setDirtyCanvas(){}
|
||||
serialize(){return {widgets_values:this.widgets.map(widget=>widget.value),properties:{...this.properties}};}
|
||||
configure(data){this.widgets.forEach((widget,index)=>{widget.value=data.widgets_values[index];});this.properties={...data.properties};}
|
||||
}
|
||||
await fixtureApp.extension.beforeRegisterNodeDef(Node,{name:'PromptManagerFocused'},fixtureApp);
|
||||
window.node=new Node();node.onNodeCreated();
|
||||
});
|
||||
await expect(page.locator('.pm-focused')).toBeVisible();
|
||||
});
|
||||
test.afterEach(async({page,browserName})=>{await saveCoverage(page,browserName);expect(page.focusedErrors).toEqual([]);});
|
||||
|
||||
test('Focused library autocompletes filters and searches without changing editor metadata', async ({page}) => {
|
||||
const requests=[];
|
||||
await page.route('**/prompt_manager/categories',route=>route.fulfill({json:{success:true,categories:['Portrait']}}));
|
||||
await page.route('**/prompt_manager/tags/suggest?**',route=>route.fulfill({json:{success:true,suggestions:[{name:'studio'}]}}));
|
||||
await page.route('**/prompt_manager/search?**',route=>{
|
||||
requests.push(new URL(route.request().url()));
|
||||
return route.fulfill({json:{success:true,results:[{text:'Filtered portrait',category:'Portrait',tags:['studio']}]}});
|
||||
});
|
||||
await page.getByLabel('Main prompt',{exact:true}).fill('My draft');
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
const category=page.getByLabel('Filter category',{exact:true});
|
||||
await category.fill('por');
|
||||
await expect(page.getByRole('option',{name:'Portrait',exact:true})).toBeVisible();
|
||||
await category.press('ArrowDown');
|
||||
await category.press('Enter');
|
||||
expect(requests).toHaveLength(0);
|
||||
const tags=page.getByLabel('Filter tags',{exact:true});
|
||||
await tags.fill('stu');
|
||||
await page.getByRole('option',{name:'studio',exact:true}).click();
|
||||
await expect(tags).toHaveValue('studio, ');
|
||||
await page.getByRole('searchbox',{name:'Search saved prompts'}).fill('portrait tag:studio tag:"film grain"');
|
||||
await page.getByRole('button',{name:'Search',exact:true}).click();
|
||||
await expect(page.locator('.pm-editor-result')).toContainText('Filtered portrait');
|
||||
expect(requests[0].searchParams.get('text')).toBe('portrait');
|
||||
expect(requests[0].searchParams.get('category')).toBe('Portrait');
|
||||
expect(requests[0].searchParams.get('tags')).toBe('studio,film grain');
|
||||
expect(await page.evaluate(()=>node.serialize().widgets_values.slice(0,3))).toEqual(['My draft','','']);
|
||||
await page.locator('.pm-editor-result').click();
|
||||
await category.fill('Landscape');
|
||||
await expect(page.locator('.pm-editor-selection button')).toBeDisabled();
|
||||
const saved=await page.evaluate(()=>{const saved=node.serialize();node.configure(saved);return saved;});
|
||||
expect(saved.properties.focusedUI.libraryFilters).toEqual({category:'Landscape',tags:'studio, '});
|
||||
await expect(category).toHaveValue('Landscape');
|
||||
await expect(tags).toHaveValue('studio, ');
|
||||
let recentRequest;
|
||||
await page.route('**/prompt_manager/recent?**',route=>{recentRequest=new URL(route.request().url());return route.fulfill({json:{success:true,results:[]}});});
|
||||
await page.getByRole('button',{name:'Recent',exact:true}).click();
|
||||
await expect.poll(()=>Boolean(recentRequest)).toBe(true);
|
||||
expect(recentRequest.searchParams.has('category')).toBe(false);
|
||||
expect(recentRequest.searchParams.has('tags')).toBe(false);
|
||||
await page.getByRole('button',{name:'Hide library',exact:true}).click();
|
||||
await expect(page.getByRole('listbox')).toBeHidden();
|
||||
});
|
||||
|
||||
test('Focused library can search by filters alone and Enter searches after autocomplete closes', async ({page}) => {
|
||||
const requests=[];
|
||||
await page.route('**/prompt_manager/search?**',route=>{requests.push(new URL(route.request().url()));return route.fulfill({json:{success:true,results:[]}});});
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
await page.getByLabel('Filter category',{exact:true}).fill('New category');
|
||||
await page.getByLabel('Filter category',{exact:true}).press('Enter');
|
||||
await expect.poll(()=>requests.length).toBe(1);
|
||||
expect(requests[0].searchParams.get('category')).toBe('New category');
|
||||
await page.getByLabel('Filter category',{exact:true}).fill('');
|
||||
await page.getByLabel('Filter tags',{exact:true}).fill('custom, ');
|
||||
await page.getByLabel('Filter tags',{exact:true}).press('Enter');
|
||||
await expect.poll(()=>requests.length).toBe(2);
|
||||
expect(requests[1].searchParams.get('tags')).toBe('custom');
|
||||
});
|
||||
|
||||
test('Focused completes saved tags with keyboard and pointer and appends commas', async ({page}) => {
|
||||
await page.route('**/prompt_manager/tags/suggest?**', route => {
|
||||
const q = new URL(route.request().url()).searchParams.get('q');
|
||||
return route.fulfill({json:{success:true,suggestions:(q === 'bl' ? ['blue eyes','blue sky'] : ['portrait','blue eyes']).map(name => ({name}))}});
|
||||
});
|
||||
await page.getByText('Metadata · category / tags',{exact:true}).click();
|
||||
const tags = page.getByLabel('Tags',{exact:true});
|
||||
await tags.fill('portrait, bl');
|
||||
await expect(page.getByRole('option',{name:'blue eyes',exact:true})).toBeVisible();
|
||||
await tags.press('ArrowDown');
|
||||
await tags.press('Enter');
|
||||
await expect(tags).toHaveValue('portrait, blue eyes, ');
|
||||
expect(await page.evaluate(()=>node.serialize().widgets_values[2])).toBe('portrait, blue eyes, ');
|
||||
await tags.fill('bl');
|
||||
await page.getByRole('option',{name:'blue sky',exact:true}).click();
|
||||
await expect(tags).toHaveValue('blue sky, ');
|
||||
await tags.fill('portrait, ');
|
||||
await expect(page.getByRole('option',{name:'portrait',exact:true})).toHaveCount(0);
|
||||
await expect(page.getByRole('option',{name:'blue eyes',exact:true})).toBeVisible();
|
||||
await tags.press('Escape');
|
||||
await expect(page.getByRole('listbox')).toBeHidden();
|
||||
await tags.fill('bl');
|
||||
await expect(page.getByRole('option',{name:'blue eyes',exact:true})).toBeVisible();
|
||||
await page.locator('.pm-focused').dispatchEvent('scroll');
|
||||
await expect(page.getByRole('listbox')).toBeHidden();
|
||||
await tags.fill('A brand new tag');
|
||||
await tags.press('Tab');
|
||||
await expect(tags).toHaveValue('A brand new tag');
|
||||
});
|
||||
|
||||
test('Focused metadata autocomplete ignores stale responses, recovers from errors and cleans up', async ({page}) => {
|
||||
let release;
|
||||
let mode='delayed';
|
||||
await page.route('**/prompt_manager/tags/suggest?**', async route => {
|
||||
const q = new URL(route.request().url()).searchParams.get('q');
|
||||
if(q === 'old') await new Promise(resolve=>{release=resolve;});
|
||||
if(mode === 'error') return route.fulfill({status:503,body:'Unavailable'});
|
||||
await route.fulfill({json:{success:true,suggestions:[{name:q+' result'}]}}).catch(()=>{});
|
||||
});
|
||||
await page.getByText('Metadata · category / tags',{exact:true}).click();
|
||||
const tags=page.getByLabel('Tags',{exact:true});
|
||||
await tags.fill('old');
|
||||
await expect.poll(()=>Boolean(release)).toBe(true);
|
||||
await tags.fill('new');
|
||||
await expect(page.getByRole('option',{name:'new result',exact:true})).toBeVisible();
|
||||
release();
|
||||
await expect(page.getByRole('option',{name:'old result',exact:true})).toHaveCount(0);
|
||||
mode='error';
|
||||
await tags.fill('unavailable');
|
||||
await expect(page.getByRole('listbox')).toBeHidden();
|
||||
await expect(tags).toHaveValue('unavailable');
|
||||
mode='ok';
|
||||
await tags.fill('retry');
|
||||
await expect(page.getByRole('option',{name:'retry result',exact:true})).toBeVisible();
|
||||
await page.evaluate(()=>node.onRemoved());
|
||||
await expect(page.getByRole('listbox')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('Focused completes one category, allows new values and dismisses suggestions on blur', async ({page}) => {
|
||||
await page.route('**/prompt_manager/categories',route=>route.fulfill({json:{success:true,categories:['Portrait','Landscape','Portrait studio']}}));
|
||||
await page.getByText('Metadata · category / tags',{exact:true}).click();
|
||||
const category=page.getByLabel('Category',{exact:true});
|
||||
await category.fill('por');
|
||||
await expect(page.getByRole('listbox',{name:'Saved categories'})).toBeVisible();
|
||||
await category.press('ArrowUp');
|
||||
await expect(page.getByRole('option',{name:'Portrait studio',exact:true})).toHaveAttribute('aria-selected','true');
|
||||
await category.press('Enter');
|
||||
await expect(category).toHaveValue('Portrait studio');
|
||||
expect(await page.evaluate(()=>node.serialize().widgets_values[1])).toBe('Portrait studio');
|
||||
await category.fill('por');
|
||||
await expect(page.getByRole('option',{name:'Portrait',exact:true})).toBeVisible();
|
||||
await category.dispatchEvent('keydown',{key:'ArrowDown',isComposing:true});
|
||||
await expect(category).not.toHaveAttribute('aria-activedescendant',/.+/);
|
||||
await page.getByLabel('Tags',{exact:true}).focus();
|
||||
await expect(page.getByRole('listbox',{name:'Saved categories'})).toBeHidden();
|
||||
await category.fill('New category');
|
||||
await expect(page.getByRole('listbox',{name:'Saved categories'})).toBeHidden();
|
||||
await expect(category).toHaveValue('New category');
|
||||
});
|
||||
|
||||
test('Focused shares suggestion controls safely with another bundled toolkit copy', async ({page}) => {
|
||||
await page.evaluate(async()=>{
|
||||
const {createScope}=await import('/prompt_manager/js/shared/toolkit/core.js');
|
||||
const first=await import('/prompt_manager/js/shared/toolkit/suggestions.js?pack=first');
|
||||
const second=await import('/prompt_manager/js/shared/toolkit/suggestions.js?pack=second');
|
||||
window.completionFixture={scope:createScope(),selected:[],inserted:[]};
|
||||
const one=first.createSuggestionView(completionFixture.scope,null);
|
||||
const two=second.createSuggestionView(completionFixture.scope,null,{label:'Other pack'});
|
||||
completionFixture.views=[one,two];
|
||||
one.hide();
|
||||
const items=[{name:'plain'},{name:'tag',type:'tag'},{name:'model',type:'lora'},{name:'nested',type:'embedding',fullPath:'folder/nested'},{name:'flat',type:'embedding',fullPath:'flat'}];
|
||||
one.render(items,index=>{completionFixture.selected.push(index);one.select(index);},item=>completionFixture.inserted.push(item.name));
|
||||
one.select(-1);
|
||||
one.position({left:innerWidth-5,bottom:innerHeight,top:innerHeight-20});
|
||||
two.render([{name:'Other pack suggestion'}],()=>{},()=>{});
|
||||
two.position({left:0,bottom:5,top:0});
|
||||
});
|
||||
const lists=page.getByRole('listbox');
|
||||
const ids=await lists.evaluateAll(elements=>elements.map(element=>element.id));
|
||||
expect(new Set(ids).size).toBe(ids.length);
|
||||
await page.getByRole('option',{name:'[LORA] model'}).hover();
|
||||
await expect(page.getByRole('option',{name:'[LORA] model'})).toHaveAttribute('aria-selected','true');
|
||||
await page.getByRole('option',{name:'[EMB] folder/nested'}).click();
|
||||
expect(await page.evaluate(()=>completionFixture.inserted)).toEqual(['nested']);
|
||||
await page.evaluate(()=>completionFixture.scope.dispose());
|
||||
await expect(lists).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('Focused previews before loading, keeps metadata optional and composes text',async({page})=>{
|
||||
await page.route('**/prompt_manager/recent?**',route=>route.fulfill({json:{success:true,results:[{id:1,text:'Library portrait',category:'Portrait',tags:['studio']}]}}));
|
||||
await page.getByRole('textbox',{name:'Main prompt',exact:true}).fill('My draft');
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
await page.getByRole('button',{name:'Recent',exact:true}).click();
|
||||
await page.getByRole('button',{name:/Library portrait/}).click();
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toHaveValue('My draft');
|
||||
await page.getByRole('switch',{name:'Load metadata'}).uncheck();
|
||||
await page.getByRole('button',{name:'Use prompt',exact:true}).click();
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toHaveValue('Library portrait');
|
||||
await page.getByText('Composition · prepend / append',{exact:true}).click();
|
||||
await page.getByLabel('Prepend',{exact:true}).fill('trigger');
|
||||
await page.getByLabel('Append',{exact:true}).fill('detailed');
|
||||
await page.getByText('Composed text preview',{exact:true}).click();
|
||||
await expect(page.locator('[data-composed-preview]')).toHaveText('trigger Library portrait detailed');
|
||||
expect(await page.evaluate(()=>node.serialize().widgets_values.slice(0,6))).toEqual(['Library portrait','','','','trigger','detailed']);
|
||||
});
|
||||
|
||||
test('Focused restores workflow values and disclosures and protects connected text',async({page})=>{
|
||||
await page.getByRole('textbox',{name:'Main prompt',exact:true}).fill('Saved draft');
|
||||
await page.getByText('Metadata · category / tags',{exact:true}).click();
|
||||
await page.getByLabel('Category',{exact:true}).fill('Saved category');
|
||||
await page.evaluate(()=>{const saved=node.serialize();node.configure(saved);});
|
||||
await expect(page.getByLabel('Category',{exact:true})).toBeVisible();
|
||||
await expect(page.getByLabel('Category',{exact:true})).toHaveValue('Saved category');
|
||||
await page.evaluate(()=>{node.inputs.push({name:'text',link:9});node.onConnectionsChange();});
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toBeHidden();
|
||||
await expect(page.getByText('Text is supplied by an upstream connection.',{exact:true})).toBeVisible();
|
||||
await page.evaluate(()=>{node.inputs.find(input=>input.name==='text').link=null;node.onConnectionsChange();});
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toHaveValue('Saved draft');
|
||||
});
|
||||
|
||||
test('Focused searches quoted tags with Enter and explicitly loads metadata', async ({page}) => {
|
||||
let requested;
|
||||
await page.route('**/prompt_manager/search?**', async route => {
|
||||
requested = new URL(route.request().url());
|
||||
await route.fulfill({json:{success:true,results:[{id:7,text:'Selected fox',category:'Animals',tags:['film grain','fox']}]}});
|
||||
});
|
||||
await page.getByRole('textbox',{name:'Main prompt',exact:true}).fill('Preserved draft');
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
const query=page.getByRole('searchbox',{name:'Search saved prompts'});
|
||||
await query.fill('fox tag:"film grain"');
|
||||
await query.press('Enter');
|
||||
await page.getByRole('button',{name:/Selected fox/}).click();
|
||||
expect(requested.searchParams.get('text')).toBe('fox');
|
||||
expect(requested.searchParams.get('tags')).toBe('film grain');
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toHaveValue('Preserved draft');
|
||||
await page.getByRole('button',{name:'Use prompt',exact:true}).click();
|
||||
await page.getByText('Metadata · category / tags',{exact:true}).click();
|
||||
await expect(page.getByLabel('Category',{exact:true})).toHaveValue('Animals');
|
||||
await expect(page.getByLabel('Tags',{exact:true})).toHaveValue('film grain, fox');
|
||||
expect(await page.evaluate(()=>node.serialize().widgets_values.slice(0,4))).toEqual(['Selected fox','Animals','film grain, fox','fox tag:"film grain"']);
|
||||
});
|
||||
|
||||
test('Focused reports empty queries, no results, and safely renders backend and HTTP errors', async ({page}) => {
|
||||
let requests=0;
|
||||
let mode='empty';
|
||||
const malicious='<img src=x onerror="window.focusedInjected=true">';
|
||||
await page.route('**/prompt_manager/search?**', async route => {
|
||||
requests++;
|
||||
if(mode==='http') return route.fulfill({status:503,body:'Unavailable'});
|
||||
await route.fulfill({json:mode==='error'?{success:false,error:malicious}:{success:true,results:[]}});
|
||||
});
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
const search=page.getByRole('button',{name:'Search',exact:true});
|
||||
await search.click();
|
||||
await expect(page.locator('.pm-editor-library [role="status"]')).toHaveText('Enter search text, a category or tags.');
|
||||
expect(requests).toBe(0);
|
||||
await page.getByRole('searchbox',{name:'Search saved prompts'}).fill('missing');
|
||||
await search.click();
|
||||
await expect(page.locator('.pm-editor-library [role="status"]')).toHaveText('No matching prompts.');
|
||||
await expect(page.locator('.pm-editor-result')).toHaveCount(0);
|
||||
mode='error';
|
||||
await search.click();
|
||||
await expect(page.getByRole('alert')).toHaveText(`Could not load prompts: ${malicious}`);
|
||||
await expect(page.locator('.pm-editor-library img')).toHaveCount(0);
|
||||
expect(await page.evaluate(()=>window.focusedInjected)).toBeUndefined();
|
||||
mode='http';
|
||||
await search.click();
|
||||
await expect(page.getByRole('alert')).toContainText('HTTP 503');
|
||||
await page.getByRole('searchbox',{name:'Search saved prompts'}).fill('');
|
||||
await search.click();
|
||||
await expect(page.getByRole('alert')).toHaveCount(0);
|
||||
await expect(page.locator('.pm-editor-library [role="status"]')).toHaveText('Enter search text, a category or tags.');
|
||||
mode='empty';
|
||||
await page.getByRole('searchbox',{name:'Search saved prompts'}).fill('missing');
|
||||
await search.click();
|
||||
await expect(page.getByRole('alert')).toHaveCount(0);
|
||||
await expect(page.locator('.pm-editor-library [role="status"]')).toHaveText('No matching prompts.');
|
||||
});
|
||||
|
||||
test('Focused treats saved prompt content and metadata as text', async ({page}) => {
|
||||
const malicious='<img src=x onerror="window.focusedInjected=true">';
|
||||
await page.route('**/prompt_manager/recent?**',route=>route.fulfill({json:{success:true,results:[{id:8,text:malicious,category:'<script>bad()</script>',tags:['<svg onload=bad()>']}]}}));
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
await page.getByRole('button',{name:'Recent',exact:true}).click();
|
||||
await page.locator('.pm-editor-result').click();
|
||||
await expect(page.locator('.pm-editor-selection .pm-editor-preview')).toHaveText(malicious);
|
||||
await expect(page.locator('.pm-editor-library img, .pm-editor-library svg, .pm-editor-library script')).toHaveCount(0);
|
||||
await page.getByRole('button',{name:'Use prompt',exact:true}).click();
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toHaveValue(malicious);
|
||||
expect(await page.evaluate(()=>window.focusedInjected)).toBeUndefined();
|
||||
});
|
||||
|
||||
async function deferLibraryRequests(page) {
|
||||
await page.evaluate(()=>{
|
||||
const originalFetch=window.fetch;
|
||||
window.libraryRequests=[];
|
||||
window.fetch=(url,options)=>{
|
||||
if(!/\/prompt_manager\/(search|recent)\?/.test(String(url))) return originalFetch(url,options);
|
||||
return new Promise(resolve=>window.libraryRequests.push({url:String(url),signal:options.signal,resolve}));
|
||||
};
|
||||
window.completeLibraryRequest=(index,text)=>libraryRequests[index].resolve(new Response(JSON.stringify({success:true,results:[{id:index,text}]}),{status:200,headers:{'Content-Type':'application/json'}}));
|
||||
});
|
||||
}
|
||||
|
||||
test('Focused keeps the latest result when searches and recent requests overlap', async ({page}) => {
|
||||
await deferLibraryRequests(page);
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
await page.getByRole('searchbox',{name:'Search saved prompts'}).fill('older search');
|
||||
await page.getByRole('button',{name:'Search',exact:true}).click();
|
||||
await page.getByRole('button',{name:'Recent',exact:true}).click();
|
||||
expect(await page.evaluate(()=>libraryRequests[0].signal.aborted)).toBe(true);
|
||||
await page.evaluate(()=>completeLibraryRequest(1,'Latest recent'));
|
||||
await expect(page.locator('.pm-editor-result')).toHaveText('Latest recent');
|
||||
await page.evaluate(()=>completeLibraryRequest(0,'Stale search'));
|
||||
await expect(page.locator('.pm-editor-result')).toHaveText('Latest recent');
|
||||
await page.getByRole('button',{name:'Recent',exact:true}).click();
|
||||
await page.getByRole('searchbox',{name:'Search saved prompts'}).fill('latest search');
|
||||
await page.getByRole('button',{name:'Search',exact:true}).click();
|
||||
expect(await page.evaluate(()=>libraryRequests[2].signal.aborted)).toBe(true);
|
||||
await page.evaluate(()=>completeLibraryRequest(3,'Latest search'));
|
||||
await expect(page.locator('.pm-editor-result')).toHaveText('Latest search');
|
||||
await page.evaluate(()=>completeLibraryRequest(2,'Stale recent'));
|
||||
await expect(page.locator('.pm-editor-result')).toHaveText('Latest search');
|
||||
});
|
||||
|
||||
test('Focused removal aborts work and stops theme and control listeners', async ({page}) => {
|
||||
await deferLibraryRequests(page);
|
||||
await page.getByRole('button',{name:'Browse library',exact:true}).click();
|
||||
await page.getByRole('button',{name:'Recent',exact:true}).click();
|
||||
await expect(page.locator('.pm-focused')).toHaveAttribute('data-ca-theme','dark');
|
||||
const removed=await page.evaluate(()=>{
|
||||
node.onRemoved();
|
||||
document.body.style.setProperty('--bg-color','#ffffff');
|
||||
fixtureApp.ui.settings.dispatchEvent(new Event('Comfy.ColorPalette.change'));
|
||||
return {requestAborted:libraryRequests[0].signal.aborted,scopeAborted:node._focusedScope.signal.aborted};
|
||||
});
|
||||
expect(removed).toEqual({requestAborted:true,scopeAborted:true});
|
||||
await page.evaluate(()=>completeLibraryRequest(0,'Late result after removal'));
|
||||
await expect(page.locator('.pm-editor-result')).toHaveCount(0);
|
||||
await expect(page.locator('.pm-focused')).toHaveAttribute('data-ca-theme','dark');
|
||||
await page.getByRole('button',{name:'Hide library',exact:true}).click();
|
||||
await expect(page.getByRole('button',{name:'Hide library',exact:true})).toBeVisible();
|
||||
});
|
||||
|
||||
test('Focused fields reflow with node width and follow live host themes', async ({page}) => {
|
||||
await page.setViewportSize({width:1000,height:1000});
|
||||
await page.getByText('Composition · prepend / append',{exact:true}).click();
|
||||
await page.evaluate(()=>{node._focusedView.element.style.width='700px';node._focusedView.element.style.height='850px';});
|
||||
const prepend=page.getByLabel('Prepend',{exact:true});
|
||||
const append=page.getByLabel('Append',{exact:true});
|
||||
let first=await prepend.boundingBox(),second=await append.boundingBox();
|
||||
expect(Math.abs(first.y-second.y)).toBeLessThan(2);
|
||||
expect(second.x).toBeGreaterThan(first.x+first.width);
|
||||
await page.evaluate(()=>{node._focusedView.element.style.width='300px';});
|
||||
await expect.poll(async()=>{const a=await prepend.boundingBox(),b=await append.boundingBox();return b.y>a.y+a.height;}).toBe(true);
|
||||
expect(await page.locator('.pm-focused').evaluate(element=>element.scrollWidth<=element.clientWidth+1)).toBe(true);
|
||||
const dark=await prepend.evaluate(element=>getComputedStyle(element).backgroundColor);
|
||||
await page.evaluate(()=>document.body.style.setProperty('--bg-color','#ffffff'));
|
||||
await expect(page.locator('.pm-focused')).toHaveAttribute('data-ca-theme','light');
|
||||
await expect.poll(()=>prepend.evaluate(element=>getComputedStyle(element).backgroundColor)).not.toBe(dark);
|
||||
await page.evaluate(()=>document.body.style.setProperty('--bg-color','#222222'));
|
||||
await expect(page.locator('.pm-focused')).toHaveAttribute('data-ca-theme','dark');
|
||||
await expect.poll(()=>prepend.evaluate(element=>getComputedStyle(element).backgroundColor)).toBe(dark);
|
||||
});
|
||||
|
||||
test('Focused composed preview preserves main prompt whitespace like queued execution', async ({page}) => {
|
||||
await page.getByRole('textbox',{name:'Main prompt',exact:true}).fill(' draft ');
|
||||
await page.getByText('Composition · prepend / append',{exact:true}).click();
|
||||
await page.getByLabel('Prepend',{exact:true}).fill(' trigger ');
|
||||
await page.getByLabel('Append',{exact:true}).fill(' detail ');
|
||||
await page.getByText('Composed text preview',{exact:true}).click();
|
||||
expect(await page.locator('[data-composed-preview]').textContent()).toBe('trigger draft detail');
|
||||
});
|
||||
|
||||
test('Focused exposes the visible text element for autocomplete and host extension listeners', async ({page}) => {
|
||||
expect(await page.evaluate(()=>node.widgets.find(widget=>widget.name==='text').inputEl===document.querySelector('.pm-focused textarea'))).toBe(true);
|
||||
});
|
||||
|
||||
test('Focused expands its main editor when the node gets taller',async({page})=>{
|
||||
await page.locator('.pm-focused').evaluate(element=>{element.style.width='550px';element.style.height='800px';});
|
||||
const editor=await page.getByRole('textbox',{name:'Main prompt',exact:true}).boundingBox();
|
||||
expect(editor.height).toBeGreaterThan(400);
|
||||
});
|
||||
|
||||
test('Focused respects the existing Web UI display preference',async({page})=>{
|
||||
await page.route('**/prompt_manager/settings',route=>route.fulfill({json:{success:true,settings:{webui_display_mode:'newtab'}}}));
|
||||
await page.evaluate(()=>{
|
||||
const Node=node.constructor;node.onRemoved();node._focusedView.element.remove();node=new Node();node.onNodeCreated();
|
||||
window.open=(...args)=>{window.openArgs=args;return {};};
|
||||
});
|
||||
await page.evaluate(()=>node._focusedSettingsReady);
|
||||
await page.getByRole('button',{name:'Web UI ↗',exact:true}).click();
|
||||
expect(await page.evaluate(()=>window.openArgs)).toEqual(['/prompt_manager/web','_blank']);
|
||||
});
|
||||
|
||||
test('Focused reports blocked Web UI popups without leaving the workflow',async({page})=>{
|
||||
await page.evaluate(()=>{window.open=()=>null;});
|
||||
await page.getByRole('button',{name:'Web UI ↗',exact:true}).click();
|
||||
await expect(page.locator('.pm-focused [role="alert"]')).toContainText('Allow popups');
|
||||
await expect(page.locator('.pm-focused')).toBeVisible();
|
||||
});
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Serve production pages and API routes with disposable state for browser tests."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import types
|
||||
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
|
||||
class TestModel:
|
||||
default_prompt = "Describe the image"
|
||||
wd14_general_threshold = 0.35
|
||||
wd14_character_threshold = 0.85
|
||||
custom_prompt = ""
|
||||
loaded = None
|
||||
|
||||
def get_models_status(self) -> dict:
|
||||
return {
|
||||
name: {"downloaded": True, "loaded": self.loaded == name}
|
||||
for name in ("gguf", "wd14-vit", "wd14-convnext", "wd14-swinv2")
|
||||
}
|
||||
|
||||
def is_model_loaded(self) -> bool:
|
||||
return self.loaded is not None
|
||||
|
||||
def get_loaded_model_type(self) -> str | None:
|
||||
return self.loaded
|
||||
|
||||
def load_model(self, model_type: str, use_gpu: bool = True) -> bool:
|
||||
self.loaded = model_type
|
||||
return True
|
||||
|
||||
def unload_model(self) -> None:
|
||||
self.loaded = None
|
||||
|
||||
def generate_tags(self, image_path: str, **kwargs: object) -> list[str]:
|
||||
return ["generated", Path(image_path).parent.parent.name]
|
||||
|
||||
|
||||
def create_app(root: Path) -> web.Application:
|
||||
os.environ["PROMPT_MANAGER_CONFIG_PATH"] = str(root / "config.json")
|
||||
os.chdir(root)
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
server = types.ModuleType("server")
|
||||
server.PromptServer = types.SimpleNamespace(
|
||||
instance=types.SimpleNamespace(app=app, routes=routes)
|
||||
)
|
||||
sys.modules["server"] = server
|
||||
folder_paths = types.ModuleType("folder_paths")
|
||||
folder_paths.base_path = str(root)
|
||||
folder_paths.models_dir = str(root / "models")
|
||||
for name in ("output", "input", "temp", "user"):
|
||||
directory = root / name
|
||||
directory.mkdir()
|
||||
setattr(folder_paths, f"get_{name}_directory", lambda path=directory: str(path))
|
||||
sys.modules["folder_paths"] = folder_paths
|
||||
|
||||
from py.config import GalleryConfig, PromptManagerConfig
|
||||
|
||||
PromptManagerConfig.DEFAULT_DB_PATH = str(root / "browser.db")
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(root / "output"), str(root / "input")]
|
||||
|
||||
from py.api import PromptManagerAPI
|
||||
from py import autotag
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
model = TestModel()
|
||||
autotag.get_autotag_service = lambda: model
|
||||
api = PromptManagerAPI()
|
||||
api.add_routes(routes)
|
||||
|
||||
def seed(extra_images: int = 0) -> None:
|
||||
with api.db.model.get_connection() as conn:
|
||||
for table in ("generated_images", "prompt_tags", "tags", "prompts"):
|
||||
conn.execute(f"DELETE FROM {table}")
|
||||
for name in ("output", "input"):
|
||||
shutil.rmtree(root / name)
|
||||
(root / name).mkdir()
|
||||
for index, (directory, text, tag) in enumerate(
|
||||
[("output", "Cat portrait", "cats"), ("input", "Dog landscape", "dogs")]
|
||||
):
|
||||
image_path = root / directory / "shared" / "same.png"
|
||||
image_path.parent.mkdir(exist_ok=True)
|
||||
info = PngInfo()
|
||||
info.add_text(
|
||||
"prompt",
|
||||
json.dumps(
|
||||
{"1": {"class_type": "CLIPTextEncode", "inputs": {"text": text}}}
|
||||
),
|
||||
)
|
||||
with Image.new("RGB", (64, 64), (index * 255, 64, 128)) as image:
|
||||
image.save(image_path, pnginfo=info)
|
||||
stamp = time.time() - index * 10
|
||||
os.utime(image_path, (stamp, stamp))
|
||||
pid = api.db.save_prompt(
|
||||
text, tags=[tag], prompt_hash=generate_prompt_hash(text)
|
||||
)
|
||||
api.db.link_image_to_prompt(pid, str(image_path))
|
||||
if index == 0:
|
||||
for number in range(extra_images):
|
||||
extra = root / directory / "bulk" / f"{number:04}.png"
|
||||
extra.parent.mkdir(exist_ok=True)
|
||||
shutil.copyfile(image_path, extra)
|
||||
api.db.link_image_to_prompt(pid, str(extra))
|
||||
api._generate_gallery_thumbnails_sync(api._get_all_output_dirs(), (150, 150))
|
||||
api.invalidate_gallery_cache()
|
||||
|
||||
@routes.post("/__test__/reset")
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
data = await request.json() if request.can_read_body else {}
|
||||
extra = data.get("extra_images", 0)
|
||||
if (
|
||||
isinstance(extra, bool)
|
||||
or not isinstance(extra, int)
|
||||
or not 0 <= extra <= 200
|
||||
):
|
||||
raise web.HTTPBadRequest()
|
||||
await api._run_in_executor(seed, extra)
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.get("/api/settings/Comfy.ColorPalette")
|
||||
async def palette(request: web.Request) -> web.Response:
|
||||
return web.json_response("dark")
|
||||
|
||||
async def cleanup(application: web.Application) -> None:
|
||||
api.db.close_all()
|
||||
|
||||
seed()
|
||||
app.on_cleanup.append(cleanup)
|
||||
app.add_routes(routes)
|
||||
return app
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--port", type=int, default=8765)
|
||||
args = parser.parse_args()
|
||||
with tempfile.TemporaryDirectory(prefix="pm-browser-") as directory:
|
||||
web.run_app(
|
||||
create_app(Path(directory).resolve()), host="127.0.0.1", port=args.port
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,344 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
|
||||
test.beforeEach(async ({ request, page, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
page.workflowErrors = [];
|
||||
page.on('pageerror', error => page.workflowErrors.push(error.message));
|
||||
page.on('response', response => {
|
||||
if (response.status() >= 500 && response.url().includes('/prompt_manager/')) {
|
||||
page.workflowErrors.push(`${response.status()} ${response.url()}`);
|
||||
}
|
||||
});
|
||||
await page.addInitScript(() => { localStorage.setItem('pm_update_notice_dismissed', '3.2.0'); });
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
expect(page.workflowErrors).toEqual([]);
|
||||
});
|
||||
|
||||
test('actual gallery form saves a prompt that the admin page can edit and reload', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await page.locator('#addPromptBtn').click();
|
||||
await page.locator('#addPromptText').fill('Browser saved landscape');
|
||||
await page.locator('#addPromptCategory').fill('Browser tests');
|
||||
await page.locator('#addPromptTagInput').fill('browser-tag');
|
||||
await page.locator('#addPromptTagInput').press('Enter');
|
||||
await page.locator('#saveNewPromptBtn').click();
|
||||
await expect(page.locator('#addPromptModal')).toBeHidden();
|
||||
await page.goto('/prompt_manager/admin');
|
||||
const text = page.locator('.prompt-text').filter({ hasText: 'Browser saved landscape' });
|
||||
await expect(text).toBeVisible();
|
||||
const id = await text.getAttribute('data-id');
|
||||
await page.locator(`[data-action="edit"][data-prompt-id="${id}"]`).click();
|
||||
const editor = page.locator(`.prompt-text[data-id="${id}"]`);
|
||||
await editor.fill('Browser edited portrait');
|
||||
await editor.press('Enter');
|
||||
await expect(page.locator('.prompt-text').filter({ hasText: 'Browser edited portrait' })).toBeVisible();
|
||||
await page.reload();
|
||||
await expect(page.locator('.prompt-text').filter({ hasText: 'Browser edited portrait' })).toBeVisible();
|
||||
const saved = await request.post('/prompt_manager/save', { data: { text: 'Browser saved landscape' } });
|
||||
expect((await saved.json()).prompt_id).not.toBe(Number(id));
|
||||
});
|
||||
|
||||
test('actual gallery review generates and persists tags for the selected root', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
|
||||
await page.locator('#autoTagBtn').click();
|
||||
const generated = page.waitForResponse(response => response.url().endsWith('/prompt_manager/autotag/single'));
|
||||
await page.locator('#startReviewBtn').click();
|
||||
const generation = await generated;
|
||||
expect(generation.request().postDataJSON().root_index).toBe(0);
|
||||
await expect(page.locator('#reviewTagsContainer')).toContainText('generated');
|
||||
const result = page.waitForResponse(response => response.url().endsWith('/prompt_manager/autotag/apply'));
|
||||
await page.locator('#applyReviewBtn').click();
|
||||
expect((await (await result).json()).success).toBe(true);
|
||||
const data = await (await request.get('/prompt_manager/recent?limit=20')).json();
|
||||
expect(data.results.find(prompt => prompt.text === 'Cat portrait').tags).toContain('generated');
|
||||
expect(data.results.find(prompt => prompt.text === 'Dog landscape').tags).not.toContain('generated');
|
||||
});
|
||||
|
||||
test('actual Tags page reloads the selected filter and opens its image viewer', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin#/tags');
|
||||
await page.locator('[data-tag-name="cats"]').click();
|
||||
await expect(page.locator('#tagPromptsGrid')).toContainText('Cat portrait');
|
||||
await page.locator('[data-tag-name="dogs"]').click();
|
||||
await expect(page.locator('#tagPromptsGrid')).toContainText('Dog landscape');
|
||||
await expect(page.locator('#tagPromptsGrid')).not.toContainText('Cat portrait');
|
||||
await page.locator('#tagPromptsGrid img').first().click();
|
||||
await expect(page.locator('.viewer-canvas img')).toBeVisible();
|
||||
await page.locator('.viewer-container').getByRole('button', { name: 'Close', exact: true }).click();
|
||||
await expect(page.locator('.viewer-container')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('actual gallery serves colliding filenames and shows the matching prompt', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
|
||||
const images = page.locator('#galleryGrid img[data-original]');
|
||||
const first = await images.nth(0).getAttribute('data-original');
|
||||
const second = await images.nth(1).getAttribute('data-original');
|
||||
expect(first).not.toBe(second);
|
||||
await page.locator('#galleryGrid .image-item[data-image-index="0"]').click();
|
||||
await expect(page.locator('#metadata-content')).toContainText('Cat portrait');
|
||||
await page.keyboard.press('ArrowRight');
|
||||
await expect(page.locator('#metadata-content')).toContainText('Dog landscape');
|
||||
await page.keyboard.press('Escape');
|
||||
await expect(page.locator('#metadata-sidebar')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('an edit collision is rejected by the backend and the page restores the original text', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
const original = page.locator('.prompt-text').filter({ hasText: 'Cat portrait' });
|
||||
const id = await original.getAttribute('data-id');
|
||||
await page.locator(`[data-action="edit"][data-prompt-id="${id}"]`).click();
|
||||
const editor = page.locator(`.prompt-text[data-id="${id}"]`);
|
||||
await editor.fill('Dog landscape');
|
||||
const rejected = page.waitForResponse(response => response.request().method() === 'PUT' && response.url().endsWith(`/prompts/${id}`));
|
||||
await editor.press('Enter');
|
||||
expect((await rejected).status()).toBe(409);
|
||||
await expect(editor).toHaveText('Cat portrait');
|
||||
await page.reload();
|
||||
await expect(page.locator(`.prompt-text[data-id="${id}"]`)).toHaveText('Cat portrait');
|
||||
});
|
||||
|
||||
test('a delayed real Tags response cannot replace a newer selection', async ({ page }) => {
|
||||
let release;
|
||||
const delayed = new Promise(resolve => { release = resolve; });
|
||||
let received;
|
||||
const waiting = new Promise(resolve => { received = resolve; });
|
||||
await page.route('**/prompt_manager/tags/cats/prompts?*', async route => {
|
||||
const response = await route.fetch();
|
||||
received();
|
||||
await delayed;
|
||||
await route.fulfill({ response });
|
||||
});
|
||||
await page.goto('/prompt_manager/admin#/tags');
|
||||
await page.locator('[data-tag-name="cats"]').click();
|
||||
await waiting;
|
||||
await page.locator('[data-tag-name="dogs"]').click();
|
||||
await expect(page.locator('#tagPromptsGrid')).toContainText('Dog landscape');
|
||||
const completed = page.waitForResponse(response => response.url().includes('/tags/cats/prompts?'));
|
||||
release();
|
||||
await completed;
|
||||
await expect(page.locator('#tagPromptsGrid')).not.toContainText('Cat portrait');
|
||||
});
|
||||
|
||||
test('gallery infinite scroll appends a real backend page while retaining existing cards', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { extra_images: 108 } })).toBeOK();
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(100);
|
||||
await page.locator('#settingsBtn').click();
|
||||
await page.locator('#infiniteScrollToggle').check();
|
||||
await page.locator('#saveSettingsBtn').click();
|
||||
await expect(page.locator('#settingsModal')).toBeHidden();
|
||||
const first = page.locator('#galleryGrid .image-item').first();
|
||||
await first.evaluate(element => { element.dataset.preserved = 'yes'; });
|
||||
await page.locator('#galleryScroll').evaluate(element => { element.scrollTop = element.scrollHeight; });
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(110);
|
||||
await expect(first).toHaveAttribute('data-preserved', 'yes');
|
||||
});
|
||||
|
||||
async function promptSnapshot(request) {
|
||||
const response = await request.get('/prompt_manager/recent?limit=100');
|
||||
await expect(response).toBeOK();
|
||||
return (await response.json()).results.map(({ id, text, tags, images }) => ({
|
||||
id, text, tags: [...tags].sort(),
|
||||
images: images.map(image => image.image_path).sort(),
|
||||
})).sort((a, b) => a.id - b.id);
|
||||
}
|
||||
|
||||
async function downloadBackup(page) {
|
||||
const downloading = page.waitForEvent('download');
|
||||
await page.locator('#backupBtn').click();
|
||||
const download = await downloading;
|
||||
expect(download.suggestedFilename()).toMatch(/^prompts_backup_.*\.db$/);
|
||||
expect(await download.failure()).toBeNull();
|
||||
const chunks = [];
|
||||
for await (const chunk of await download.createReadStream()) chunks.push(chunk);
|
||||
const buffer = Buffer.concat(chunks);
|
||||
expect(buffer.subarray(0, 16).toString()).toBe('SQLite format 3\0');
|
||||
return { name: download.suggestedFilename(), mimeType: 'application/octet-stream', buffer };
|
||||
}
|
||||
|
||||
async function chooseRestore(page, file) {
|
||||
await page.locator('#restoreBtn').click();
|
||||
await expect(page.locator('#restoreModal')).toBeVisible();
|
||||
await page.locator('#restoreFileInput').setInputFiles(file);
|
||||
}
|
||||
|
||||
function confirmNextDialog(page, accept, message) {
|
||||
page.once('dialog', async dialog => {
|
||||
expect(dialog.type()).toBe('confirm');
|
||||
expect(dialog.message()).toContain(message);
|
||||
await (accept ? dialog.accept() : dialog.dismiss());
|
||||
});
|
||||
}
|
||||
|
||||
async function deletePromptButton(page, text) {
|
||||
const prompt = page.locator('.prompt-text').filter({ hasText: text });
|
||||
await expect(prompt).toBeVisible();
|
||||
const id = await prompt.getAttribute('data-id');
|
||||
return page.locator(`[data-action="delete"][data-prompt-id="${id}"]`);
|
||||
}
|
||||
|
||||
test('downloaded backup restores deleted prompts, tags, and image links through the actual upload form', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await expect(page.locator('.prompt-text')).toHaveCount(2);
|
||||
const before = await promptSnapshot(request);
|
||||
expect(before.map(prompt => prompt.tags)).toEqual([['cats'], ['dogs']]);
|
||||
expect(before.every(prompt => prompt.images.length === 1)).toBe(true);
|
||||
const backup = await downloadBackup(page);
|
||||
const button = await deletePromptButton(page, 'Cat portrait');
|
||||
confirmNextDialog(page, true, 'delete this prompt');
|
||||
await button.click();
|
||||
await expect(page.locator('.prompt-text')).toHaveCount(1);
|
||||
await expect(await request.post('/prompt_manager/save', { data: { text: 'Added after backup', tags: ['temporary'] } })).toBeOK();
|
||||
expect(await promptSnapshot(request)).not.toEqual(before);
|
||||
await chooseRestore(page, backup);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
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();
|
||||
await expect(page.locator('.prompt-text').filter({ hasText: 'Added after backup' })).toHaveCount(0);
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
await page.goto('/prompt_manager/admin#/tags');
|
||||
await page.locator('[data-tag-name="cats"]').click();
|
||||
await expect(page.locator('#tagPromptsGrid')).toContainText('Cat portrait');
|
||||
await page.locator('#tagPromptsGrid img').first().click();
|
||||
await expect(page.locator('.viewer-canvas img')).toBeVisible();
|
||||
await expect.poll(() => page.locator('.viewer-canvas img').evaluate(image => image.complete && image.naturalWidth > 0)).toBe(true);
|
||||
});
|
||||
|
||||
test('restore validation and cancellation preserve live data and allow a valid retry', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
const before = await promptSnapshot(request);
|
||||
const backup = await downloadBackup(page);
|
||||
let uploads = 0;
|
||||
page.on('request', request => { if (request.url().endsWith('/prompt_manager/restore')) uploads++; });
|
||||
await page.locator('#restoreBtn').click();
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.getByText('Please select a database file to restore', { exact: true })).toBeVisible();
|
||||
await page.locator('#restoreFileInput').setInputFiles({ ...backup, name: 'backup.txt' });
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.getByText('Please select a valid .db file', { exact: true })).toBeVisible();
|
||||
expect(uploads).toBe(0);
|
||||
await page.locator('#restoreFileInput').setInputFiles(backup);
|
||||
confirmNextDialog(page, false, 'replace your current database');
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.locator('#restoreModal')).toBeVisible();
|
||||
expect(uploads).toBe(0);
|
||||
await page.locator('#restoreFileInput').setInputFiles({ name: 'corrupt.db', mimeType: 'application/octet-stream', buffer: Buffer.from('not a SQLite database') });
|
||||
confirmNextDialog(page, true, 'corrupt.db');
|
||||
const rejected = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore'));
|
||||
await page.locator('#confirmRestore').click();
|
||||
expect((await rejected).status()).toBe(400);
|
||||
await expect(page.getByText(/Restore failed:/)).toBeVisible();
|
||||
await expect(page.locator('#restoreModal')).toBeVisible();
|
||||
await expect(page.locator('#confirmRestore')).toBeEnabled();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
await page.locator('#restoreFileInput').setInputFiles(backup);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
});
|
||||
|
||||
test('backup and restore recover from interrupted network requests without losing data', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
const before = await promptSnapshot(request);
|
||||
await page.route('**/prompt_manager/backup', route => route.abort('failed'), { times: 1 });
|
||||
await page.locator('#backupBtn').click();
|
||||
await expect(page.getByText(/Backup failed:/)).toBeVisible();
|
||||
const backup = await downloadBackup(page);
|
||||
await chooseRestore(page, backup);
|
||||
await page.route('**/prompt_manager/restore', route => route.abort('failed'), { times: 1 });
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.getByText(/Restore failed:/)).toBeVisible();
|
||||
await expect(page.locator('#confirmRestore')).toBeEnabled();
|
||||
await expect(page.locator('#restoreModal')).toBeVisible();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
});
|
||||
|
||||
test('prompt deletion requires confirmation and can retry after a network failure', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
const before = await promptSnapshot(request);
|
||||
const button = await deletePromptButton(page, 'Cat portrait');
|
||||
let deletes = 0;
|
||||
page.on('request', request => { if (request.method() === 'DELETE') deletes++; });
|
||||
confirmNextDialog(page, false, 'delete this prompt');
|
||||
await button.click();
|
||||
expect(deletes).toBe(0);
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
await page.route('**/prompt_manager/delete/*', route => route.abort('failed'), { times: 1 });
|
||||
confirmNextDialog(page, true, 'delete this prompt');
|
||||
await button.click();
|
||||
await expect(page.getByText('Failed to delete prompt', { exact: true })).toBeVisible();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
confirmNextDialog(page, true, 'delete this prompt');
|
||||
await button.click();
|
||||
await expect(page.locator('.prompt-text')).toHaveCount(1);
|
||||
await page.reload();
|
||||
await expect(page.locator('.prompt-text')).toHaveCount(1);
|
||||
expect(await promptSnapshot(request)).toEqual(before.filter(prompt => prompt.text !== 'Cat portrait'));
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
|
||||
for (const url of await page.locator('#galleryGrid img[data-original]').evaluateAll(images => images.map(image => image.dataset.original))) {
|
||||
await expect(await request.get(url)).toBeOK();
|
||||
}
|
||||
});
|
||||
|
||||
test('duplicate cleanup cancels safely and deletes only the selected file after a retry', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { extra_images: 2 } })).toBeOK();
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(4);
|
||||
const originals = await page.locator('#galleryGrid img[data-original]').evaluateAll(images => images.map(image => image.dataset.original));
|
||||
await page.locator('#settingsBtn').click();
|
||||
await page.locator('#scanDuplicatesBtn').click();
|
||||
await expect(page.locator('#duplicatesCount')).toHaveText('1');
|
||||
const candidates = page.locator('.duplicate-checkbox');
|
||||
await expect(candidates).toHaveCount(2);
|
||||
const selected = await candidates.first().getAttribute('data-image-path');
|
||||
const retained = await candidates.last().getAttribute('data-image-path');
|
||||
const serve = path => `/prompt_manager/images/serve/${path.split(/[\\/]/).map(encodeURIComponent).join('/')}`;
|
||||
await candidates.first().check();
|
||||
let deletes = 0;
|
||||
page.on('request', request => { if (request.url().endsWith('/delete_duplicate_images')) deletes++; });
|
||||
confirmNextDialog(page, false, 'permanently delete 1 duplicate image files');
|
||||
await page.locator('#removeDuplicatesBtn').click();
|
||||
expect(deletes).toBe(0);
|
||||
await expect(await request.get(serve(selected))).toBeOK();
|
||||
await page.route('**/prompt_manager/delete_duplicate_images', route => route.abort('failed'), { times: 1 });
|
||||
confirmNextDialog(page, true, 'permanently delete 1 duplicate image files');
|
||||
await page.locator('#removeDuplicatesBtn').click();
|
||||
await expect(page.getByText('Failed to delete duplicate images', { exact: true })).toBeVisible();
|
||||
await expect(page.locator('#removeDuplicatesBtn')).toBeEnabled();
|
||||
await expect(await request.get(serve(selected))).toBeOK();
|
||||
confirmNextDialog(page, true, 'permanently delete 1 duplicate image files');
|
||||
const deleted = page.waitForResponse(response => response.url().endsWith('/delete_duplicate_images'));
|
||||
await page.locator('#removeDuplicatesBtn').click();
|
||||
expect((await (await deleted).json())).toMatchObject({ success: true, deleted_count: 1, failed_count: 0 });
|
||||
await expect(page.locator('#duplicatesModal')).toBeHidden();
|
||||
expect((await request.get(serve(selected))).status()).toBe(404);
|
||||
await expect(await request.get(serve(retained))).toBeOK();
|
||||
await page.reload();
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(3);
|
||||
const remaining = await page.locator('#galleryGrid img[data-original]').evaluateAll(images => images.map(image => image.dataset.original));
|
||||
expect(remaining).toHaveLength(originals.length - 1);
|
||||
for (const url of remaining) {
|
||||
expect(originals).toContain(url);
|
||||
await expect(await request.get(url)).toBeOK();
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,346 @@
|
||||
import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { installHostHooks } from './host-fixture.mjs';
|
||||
|
||||
installHostHooks(test, ['PromptManager', 'PromptManagerText', 'PromptSearchList', 'PromptManagerFocused']);
|
||||
|
||||
async function addFocused(page) {
|
||||
return page.evaluate(async () => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = LiteGraph.createNode('PromptManagerFocused');
|
||||
if (!node) throw new Error('PromptManagerFocused must be registered by the real host');
|
||||
app.graph.add(node);
|
||||
node.setSize([680, 1050]);
|
||||
app.canvas.ds.scale = .8;
|
||||
app.canvas.centerOnNode(node);
|
||||
return node.id;
|
||||
});
|
||||
}
|
||||
|
||||
async function model(page, id) {
|
||||
return page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
return Object.fromEntries(app.graph.getNodeById(id).widgets.map(widget => [widget.name, widget.value]));
|
||||
}, id);
|
||||
}
|
||||
|
||||
async function seed(request, text, category = 'Focused host tests', tags = ['focused-host']) {
|
||||
const response = await request.post('/prompt_manager/save', { data: { text, category, tags } });
|
||||
expect(response.ok()).toBe(true);
|
||||
expect((await response.json()).success).toBe(true);
|
||||
}
|
||||
|
||||
async function browse(page, text) {
|
||||
const root = page.locator('.pm-focused');
|
||||
await root.getByRole('button', { name: 'Browse library', exact: true }).click();
|
||||
await root.getByRole('button', { name: 'Recent', exact: true }).click();
|
||||
await root.locator('.pm-editor-results').getByRole('button').filter({ has: page.getByText(text, { exact: true }) }).click();
|
||||
return root;
|
||||
}
|
||||
|
||||
test('Focused definition preserves Classic encoding and exposes composition STRING sockets', async ({ request }) => {
|
||||
const definitions = await (await request.get('/object_info')).json();
|
||||
const focused = definitions.PromptManagerFocused;
|
||||
expect(focused).toBeDefined();
|
||||
expect(focused.display_name).toBe('Prompt Manager (Focused)');
|
||||
expect(focused.output).toEqual(definitions.PromptManager.output);
|
||||
expect(focused.input.required).toEqual(definitions.PromptManager.input.required);
|
||||
for (const name of ['prepend_text', 'append_text']) {
|
||||
expect(focused.input.optional[name][0]).toBe('STRING');
|
||||
expect(focused.input.optional[name][1]).toMatchObject({ forceInput: true, default: '' });
|
||||
}
|
||||
});
|
||||
|
||||
test('Focused preview, explicit load, composition, renderer changes and workflow reload preserve state', async ({ page, request }, testInfo) => {
|
||||
const text = `Focused selected ${testInfo.project.name}`;
|
||||
await seed(request, text);
|
||||
const id = await addFocused(page);
|
||||
const root = page.locator('.pm-focused');
|
||||
const editor = root.getByRole('textbox', { name: 'Main prompt', exact: true });
|
||||
await expect(root).toBeVisible();
|
||||
expect(await editor.evaluate(async (element, id) => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
return app.graph.getNodeById(id).widgets.find(widget => widget.name === 'text').inputEl === element;
|
||||
}, id)).toBe(true);
|
||||
await editor.fill('Original editing text');
|
||||
await browse(page, text);
|
||||
await expect(editor).toHaveValue('Original editing text');
|
||||
expect((await model(page, id)).text).toBe('Original editing text');
|
||||
await root.getByRole('button', { name: 'Use prompt', exact: true }).click();
|
||||
await expect(editor).toHaveValue(text);
|
||||
await expect.poll(() => model(page, id)).toMatchObject({ text, category: 'Focused host tests', tags: 'focused-host' });
|
||||
await root.getByText('Composition · prepend / append', { exact: true }).click();
|
||||
await root.getByRole('textbox', { name: 'Prepend', exact: true }).fill('LoRA trigger');
|
||||
await root.getByRole('textbox', { name: 'Append', exact: true }).fill('high detail');
|
||||
await root.getByText('Composed text preview', { exact: true }).click();
|
||||
await expect(root).toContainText(`LoRA trigger ${text} high detail`);
|
||||
for (const renderer of [!testInfo.project.metadata.initialRenderer, testInfo.project.metadata.initialRenderer]) {
|
||||
await page.evaluate(async renderer => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.VueNodes.Enabled', renderer);
|
||||
app.canvas.centerOnNode(app.graph._nodes[0]);
|
||||
}, renderer);
|
||||
await expect.poll(() => page.evaluate(() => LiteGraph.vueNodesMode)).toBe(renderer);
|
||||
await expect(root).toHaveCount(1);
|
||||
await expect(editor).toHaveValue(text);
|
||||
}
|
||||
const queued = await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const prompt = await app.graphToPrompt();
|
||||
await app.loadGraphData(app.graph.serialize());
|
||||
return prompt.output[String(id)].inputs;
|
||||
}, id);
|
||||
expect(queued).toMatchObject({ text, category: 'Focused host tests', tags: 'focused-host', prepend_text: 'LoRA trigger', append_text: 'high detail' });
|
||||
await expect(editor).toHaveValue(text);
|
||||
await expect.poll(() => model(page, id)).toMatchObject({ text, prepend_text: 'LoRA trigger', append_text: 'high detail' });
|
||||
await page.screenshot({ path: testInfo.outputPath('focused-reloaded.png') });
|
||||
await page.evaluate(async id => { const { app } = await import('/scripts/app.js'); app.graph.remove(app.graph.getNodeById(id)); }, id);
|
||||
await expect(root).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('Focused search preview respects metadata opt-out', async ({ page, request }, testInfo) => {
|
||||
const text = `Focused search ${testInfo.project.name}`;
|
||||
await seed(request, text, 'Incoming category', ['incoming-tag']);
|
||||
const id = await addFocused(page);
|
||||
const root = page.locator('.pm-focused');
|
||||
await root.getByText('Metadata · category / tags', { exact: true }).click();
|
||||
await root.getByRole('combobox', { name: 'Category', exact: true }).fill('Keep category');
|
||||
await root.getByRole('combobox', { name: 'Tags', exact: true }).fill('keep-tag');
|
||||
await root.getByRole('button', { name: 'Browse library', exact: true }).click();
|
||||
await root.getByRole('searchbox', { name: 'Search saved prompts', exact: true }).fill(text);
|
||||
await root.getByRole('button', { name: 'Search', exact: true }).click();
|
||||
await root.locator('.pm-editor-results').getByRole('button').filter({ has: page.getByText(text, { exact: true }) }).click();
|
||||
await root.getByRole('switch', { name: 'Load metadata', exact: true }).uncheck();
|
||||
await root.getByRole('button', { name: 'Use prompt', exact: true }).click();
|
||||
await expect.poll(() => model(page, id)).toMatchObject({ text, category: 'Keep category', tags: 'keep-tag' });
|
||||
});
|
||||
|
||||
test('Focused connected STRING inputs remain linked through UI actions and reload', async ({ page, request }, testInfo) => {
|
||||
const text = `Connected library ${testInfo.project.name}`;
|
||||
await seed(request, text);
|
||||
const id = await addFocused(page);
|
||||
const root = page.locator('.pm-focused');
|
||||
await root.getByRole('textbox', { name: 'Main prompt', exact: true }).fill('Local main');
|
||||
await root.getByText('Composition · prepend / append', { exact: true }).click();
|
||||
await root.getByRole('textbox', { name: 'Prepend', exact: true }).fill('Local prepend');
|
||||
await root.getByRole('textbox', { name: 'Append', exact: true }).fill('Local append');
|
||||
await root.getByText('Composition · prepend / append', { exact: true }).click();
|
||||
const links = await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const { convertToInput } = await import('/extensions/core/widgetInputs.js');
|
||||
const target = app.graph.getNodeById(id);
|
||||
convertToInput(target, target.widgets.find(widget => widget.name === 'text'));
|
||||
const result = {};
|
||||
for (const name of ['text', 'prepend_text', 'append_text']) {
|
||||
const source = LiteGraph.createNode('PrimitiveString');
|
||||
app.graph.add(source);
|
||||
source.pos = [target.pos[0] - 360, target.pos[1] + Object.keys(result).length * 150];
|
||||
source.widgets.find(widget => widget.name === 'value').value = `Connected ${name}`;
|
||||
const slot = target.inputs.findIndex(input => input.name === name);
|
||||
if (slot < 0) throw new Error(`Missing ${name} input socket`);
|
||||
source.connect(0, target, slot);
|
||||
result[name] = [String(source.id), 0];
|
||||
}
|
||||
app.canvas.centerOnNode(target);
|
||||
return result;
|
||||
}, id);
|
||||
await browse(page, text);
|
||||
await expect(root.getByRole('textbox', { name: 'Main prompt', exact: true })).toBeHidden();
|
||||
await expect(root.getByRole('button', { name: 'Use prompt', exact: true })).toBeDisabled();
|
||||
await root.getByText('Composition · prepend / append', { exact: true }).click();
|
||||
for (const name of ['Prepend', 'Append']) await expect(root.getByRole('textbox', { name, exact: true })).toBeHidden();
|
||||
const guardedValues = await model(page, id);
|
||||
await root.locator('textarea, input').evaluateAll(inputs => {
|
||||
for (const input of inputs) {
|
||||
if (!input.hidden && !input.closest('[hidden]')) continue;
|
||||
input.value = 'Attempt to replace upstream text';
|
||||
input.dispatchEvent(new Event('input', { bubbles: true }));
|
||||
}
|
||||
});
|
||||
expect(await model(page, id)).toMatchObject(guardedValues);
|
||||
for (let reload = 0; reload < 2; reload++) {
|
||||
const inputs = await page.evaluate(async ({ id, reload }) => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
if (reload) await app.loadGraphData(app.graph.serialize());
|
||||
return (await app.graphToPrompt()).output[String(id)].inputs;
|
||||
}, { id, reload });
|
||||
expect(inputs).toMatchObject(links);
|
||||
}
|
||||
await expect(root.getByRole('button', { name: 'Use prompt', exact: true, includeHidden: true })).toBeDisabled();
|
||||
const disconnected = await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = app.graph.getNodeById(id);
|
||||
for (const name of ['text', 'prepend_text', 'append_text']) node.disconnectInput(node.inputs.findIndex(input => input.name === name));
|
||||
return (await app.graphToPrompt()).output[String(id)].inputs;
|
||||
}, id);
|
||||
expect(disconnected).toMatchObject({ text: 'Local main', prepend_text: 'Local prepend', append_text: 'Local append' });
|
||||
for (const [name, value] of [['Main prompt', 'Local main'], ['Prepend', 'Local prepend'], ['Append', 'Local append']]) {
|
||||
const input = page.getByRole('textbox', { name, exact: true });
|
||||
await expect(input).toHaveCount(1);
|
||||
await expect(input).toBeEditable();
|
||||
await expect(input).toHaveValue(value);
|
||||
}
|
||||
await page.screenshot({ path: testInfo.outputPath('focused-disconnected.png') });
|
||||
});
|
||||
|
||||
test('Focused follows actual ComfyUI light and dark palettes', async ({ page }, testInfo) => {
|
||||
await addFocused(page);
|
||||
for (const palette of ['light', 'dark']) {
|
||||
await page.evaluate(async palette => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.ColorPalette', palette);
|
||||
}, palette);
|
||||
const root = page.locator('.pm-focused');
|
||||
await expect(root).toHaveAttribute('data-ca-theme', palette);
|
||||
await expect(root.getByRole('button', { name: 'Browse library', exact: true })).toHaveCSS('background-color', 'rgb(8, 119, 187)');
|
||||
await page.screenshot({ path: testInfo.outputPath(`focused-${palette}.png`) });
|
||||
}
|
||||
});
|
||||
|
||||
test('Focused Web UI button opens actual standalone pages', async ({ page }) => {
|
||||
await addFocused(page);
|
||||
const popupPromise = page.waitForEvent('popup');
|
||||
await page.locator('.pm-focused').getByRole('button', { name: 'Web UI ↗', exact: true }).click();
|
||||
const popup = await popupPromise;
|
||||
await popup.waitForLoadState('domcontentloaded');
|
||||
await expect(popup).toHaveURL(/\/prompt_manager\/(web|admin)/);
|
||||
await expect(popup.locator('body')).toContainText('Prompt');
|
||||
await expect(page.locator('.pm-focused').getByRole('status')).toHaveText('Web UI opened.');
|
||||
await popup.goto(new URL('/prompt_manager/gallery', page.url()).href);
|
||||
await expect(popup.locator('#addPromptBtn')).toBeVisible();
|
||||
await popup.close();
|
||||
});
|
||||
|
||||
test('Focused metadata completion uses saved categories and tags across renderers and reload', async ({ page, request }, testInfo) => {
|
||||
const category = `Host completion category ${testInfo.project.name}`;
|
||||
const tag = `host-completion-${testInfo.project.name}`;
|
||||
await seed(request, `Saved autocomplete ${testInfo.project.name}`, category, [tag]);
|
||||
const id = await addFocused(page);
|
||||
const root = page.locator('.pm-focused');
|
||||
await root.getByRole('textbox', { name: 'Main prompt', exact: true }).fill('Autocomplete workflow');
|
||||
await root.getByText('Metadata · category / tags', { exact: true }).click();
|
||||
const categoryInput = root.getByRole('combobox', { name: 'Category', exact: true });
|
||||
const tagInput = root.getByRole('combobox', { name: 'Tags', exact: true });
|
||||
await categoryInput.fill(category.slice(0, -3));
|
||||
const categories = page.getByRole('listbox', { name: 'Saved categories', exact: true });
|
||||
await expect(categories.getByRole('option', { name: category, exact: true })).toBeVisible();
|
||||
await categoryInput.press('ArrowDown');
|
||||
await expect(categoryInput).toHaveAttribute('aria-activedescendant', /.+/);
|
||||
await categoryInput.press('Enter');
|
||||
await expect(categoryInput).toHaveValue(category);
|
||||
await expect(categoryInput).toHaveAttribute('aria-expanded', 'false');
|
||||
for (const palette of ['light', 'dark']) {
|
||||
await page.evaluate(async palette => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.ColorPalette', palette);
|
||||
}, palette);
|
||||
await tagInput.fill(`keep, ${tag.slice(0, -3)}`);
|
||||
const tags = page.getByRole('listbox', { name: 'Saved tags', exact: true });
|
||||
const option = tags.getByRole('option', { name: tag, exact: true });
|
||||
await expect(option).toBeVisible();
|
||||
await expect(tags).toHaveAttribute('data-ca-theme', palette);
|
||||
const bounds = await tags.boundingBox();
|
||||
const viewport = page.viewportSize();
|
||||
expect(bounds.x).toBeGreaterThanOrEqual(0);
|
||||
expect(bounds.y).toBeGreaterThanOrEqual(0);
|
||||
expect(bounds.x + bounds.width).toBeLessThanOrEqual(viewport.width);
|
||||
expect(bounds.y + bounds.height).toBeLessThanOrEqual(viewport.height);
|
||||
await page.screenshot({ path: testInfo.outputPath(`focused-metadata-${palette}.png`) });
|
||||
await option.click();
|
||||
await expect(tagInput).toHaveValue(`keep, ${tag}, `);
|
||||
}
|
||||
const expected = { text: 'Autocomplete workflow', category, tags: `keep, ${tag}, ` };
|
||||
await expect.poll(() => model(page, id)).toMatchObject(expected);
|
||||
for (const renderer of [!testInfo.project.metadata.initialRenderer, testInfo.project.metadata.initialRenderer]) {
|
||||
await page.evaluate(async renderer => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.VueNodes.Enabled', renderer);
|
||||
app.canvas.centerOnNode(app.graph._nodes[0]);
|
||||
}, renderer);
|
||||
await expect.poll(() => page.evaluate(() => LiteGraph.vueNodesMode)).toBe(renderer);
|
||||
await expect(categoryInput).toHaveValue(category);
|
||||
await expect(tagInput).toHaveValue(expected.tags);
|
||||
}
|
||||
const queued = await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const prompt = await app.graphToPrompt();
|
||||
await app.loadGraphData(app.graph.serialize());
|
||||
return prompt.output[String(id)].inputs;
|
||||
}, id);
|
||||
expect(queued).toMatchObject(expected);
|
||||
await expect(categoryInput).toHaveValue(category);
|
||||
await expect(tagInput).toHaveValue(expected.tags);
|
||||
await categoryInput.fill(category.slice(0, -3));
|
||||
await expect(categories).toBeVisible();
|
||||
await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
app.graph.remove(app.graph.getNodeById(id));
|
||||
}, id);
|
||||
await expect(page.getByRole('listbox', { name: 'Saved categories', includeHidden: true })).toHaveCount(0);
|
||||
await expect(page.getByRole('listbox', { name: 'Saved tags', includeHidden: true })).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('Focused library autocomplete filters saved prompts without editing draft metadata', async ({ page, request }, testInfo) => {
|
||||
const category = `Library category ${testInfo.project.name}`;
|
||||
const tag = `library-filter-${testInfo.project.name}`;
|
||||
const matchingText = `Matching library prompt ${testInfo.project.name}`;
|
||||
await seed(request, matchingText, category, [tag]);
|
||||
await seed(request, `Wrong category ${testInfo.project.name}`, `Other category ${testInfo.project.name}`, [tag]);
|
||||
await seed(request, `Wrong tag ${testInfo.project.name}`, category, [`other-filter-${testInfo.project.name}`]);
|
||||
const id = await addFocused(page);
|
||||
const root = page.locator('.pm-focused');
|
||||
await root.getByRole('textbox', { name: 'Main prompt', exact: true }).fill('Keep my draft');
|
||||
await root.getByText('Metadata · category / tags', { exact: true }).click();
|
||||
await root.getByRole('combobox', { name: 'Category', exact: true }).fill('Draft category');
|
||||
await root.getByRole('combobox', { name: 'Tags', exact: true }).fill('draft-tag');
|
||||
await root.getByRole('button', { name: 'Browse library', exact: true }).click();
|
||||
const filterCategory = root.getByRole('combobox', { name: 'Filter category', exact: true });
|
||||
const filterTags = root.getByRole('combobox', { name: 'Filter tags', exact: true });
|
||||
await filterCategory.fill(category.slice(0, -3));
|
||||
await expect(page.getByRole('listbox', { name: 'Saved categories', exact: true }).getByRole('option', { name: category, exact: true })).toBeVisible();
|
||||
await filterCategory.press('ArrowDown');
|
||||
await filterCategory.press('Enter');
|
||||
await expect(filterCategory).toHaveValue(category);
|
||||
for (const palette of ['light', 'dark']) {
|
||||
await page.evaluate(async palette => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.ColorPalette', palette);
|
||||
}, palette);
|
||||
await filterTags.fill(tag.slice(0, -3));
|
||||
const option = page.getByRole('listbox', { name: 'Saved tags', exact: true }).getByRole('option', { name: tag, exact: true });
|
||||
await expect(option).toBeVisible();
|
||||
await page.screenshot({ path: testInfo.outputPath(`focused-library-filters-${palette}.png`) });
|
||||
await option.click();
|
||||
await expect(filterTags).toHaveValue(`${tag}, `);
|
||||
}
|
||||
const draft = { text: 'Keep my draft', category: 'Draft category', tags: 'draft-tag' };
|
||||
await root.getByRole('button', { name: 'Search', exact: true }).click();
|
||||
const results = root.locator('.pm-editor-results').getByRole('button');
|
||||
await expect(results).toHaveCount(1);
|
||||
await expect(results.first()).toContainText(matchingText);
|
||||
await results.first().click();
|
||||
await expect.poll(() => model(page, id)).toMatchObject(draft);
|
||||
for (const renderer of [!testInfo.project.metadata.initialRenderer, testInfo.project.metadata.initialRenderer]) {
|
||||
await page.evaluate(async renderer => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.VueNodes.Enabled', renderer);
|
||||
app.canvas.centerOnNode(app.graph._nodes[0]);
|
||||
}, renderer);
|
||||
await expect.poll(() => page.evaluate(() => LiteGraph.vueNodesMode)).toBe(renderer);
|
||||
await expect(filterCategory).toHaveValue(category);
|
||||
await expect(filterTags).toHaveValue(`${tag}, `);
|
||||
await expect.poll(() => model(page, id)).toMatchObject(draft);
|
||||
}
|
||||
const queued = await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const prompt = await app.graphToPrompt();
|
||||
await app.loadGraphData(app.graph.serialize());
|
||||
return prompt.output[String(id)].inputs;
|
||||
}, id);
|
||||
expect(queued).toMatchObject(draft);
|
||||
await expect(filterCategory).toHaveValue(category);
|
||||
await expect(filterTags).toHaveValue(`${tag}, `);
|
||||
await root.getByRole('button', { name: 'Search', exact: true }).click();
|
||||
await expect(results).toHaveCount(1);
|
||||
await expect(results.first()).toContainText(matchingText);
|
||||
await expect.poll(() => model(page, id)).toMatchObject(draft);
|
||||
});
|
||||
@@ -0,0 +1,43 @@
|
||||
import { expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import runtime from '../../tools/vendor/comfyui-runtime.json' with { type: 'json' };
|
||||
|
||||
export function installHostHooks(test, expectedTypes = ['PromptManager', 'PromptManagerText', 'PromptSearchList']) {
|
||||
test.beforeEach(async ({ page, context, request }, testInfo) => {
|
||||
const response = await request.get('/system_stats');
|
||||
expect(response.ok()).toBe(true);
|
||||
const info = (await response.json()).system;
|
||||
expect(info.comfyui_version).toBe(runtime.serverVersion);
|
||||
expect(info.pytorch_version).toBe(`${runtime.torchVersion}+cpu`);
|
||||
expect(info.required_frontend_version).toBe(runtime.frontendVersion);
|
||||
expect(info.comfy_package_versions.find(item => item.name === 'comfyui-frontend-package')?.installed).toBe(runtime.frontendVersion);
|
||||
await testInfo.attach('host-versions.json', { body: JSON.stringify(info, null, 2), contentType: 'application/json' });
|
||||
page.hostErrors = [];
|
||||
page.on('pageerror', error => page.hostErrors.push(error.message));
|
||||
await context.route('**/*', async route => {
|
||||
const request = route.request();
|
||||
const path = new URL(request.url()).pathname.replace(/^\/api\//, '/');
|
||||
if (path === '/settings' && request.method() === 'GET') {
|
||||
const response = await route.fetch();
|
||||
return route.fulfill({ json: { ...await response.json(),
|
||||
'Comfy.VueNodes.Enabled': testInfo.project.metadata.initialRenderer,
|
||||
'Comfy.TutorialCompleted': true, 'Comfy.Workflow.Persist': false,
|
||||
'Comfy.Notification.ShowVersionUpdates': false } });
|
||||
}
|
||||
if (!['GET', 'HEAD', 'OPTIONS'].includes(request.method()) && !path.startsWith('/prompt_manager/')) return route.fulfill({ json: {} });
|
||||
return route.continue();
|
||||
});
|
||||
await page.goto('/', { waitUntil: 'domcontentloaded' });
|
||||
await expect(page.getByRole('status', { name: 'Loading ComfyUI' })).toBeHidden({ timeout: 30000 });
|
||||
await page.waitForFunction(async expectedTypes => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
return !!app.graph && expectedTypes.every(name => globalThis.LiteGraph?.registered_node_types?.[name]);
|
||||
}, expectedTypes);
|
||||
await expect.poll(() => page.evaluate(() => LiteGraph.vueNodesMode)).toBe(testInfo.project.metadata.initialRenderer);
|
||||
await page.evaluate(async () => { const { app } = await import('/scripts/app.js'); app.graph.clear(); });
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page }) => {
|
||||
expect(page.hostErrors).toEqual([]);
|
||||
});
|
||||
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { installHostHooks } from './host-fixture.mjs';
|
||||
|
||||
installHostHooks(test);
|
||||
|
||||
for (const type of ['PromptManager', 'PromptManagerText']) {
|
||||
test(`${type} selects real database prompts and preserves text across renderers and reload`, async ({ page, request }, testInfo) => {
|
||||
const text = `Host test ${type} ${testInfo.project.name}`;
|
||||
const saved = await request.post('/prompt_manager/save', { data: { text, category: 'Host tests', tags: ['host-test'] } });
|
||||
expect(saved.ok()).toBe(true);
|
||||
expect((await saved.json()).success).toBe(true);
|
||||
const id = await page.evaluate(async type => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = LiteGraph.createNode(type);
|
||||
app.graph.add(node);
|
||||
node.setSize([650, 800]);
|
||||
app.canvas.centerOnNode(node);
|
||||
return node.id;
|
||||
}, type);
|
||||
const recent = page.getByRole('button', { name: '📋 Recent', exact: true });
|
||||
await expect(recent).toBeVisible();
|
||||
if (type === 'PromptManager') {
|
||||
await expect(page.locator('.pm-classic')).toBeVisible();
|
||||
await expect(page.getByRole('button', { name: '🔍 Search', exact: true })).toHaveCSS('background-color', 'rgb(8, 119, 187)');
|
||||
}
|
||||
await recent.click();
|
||||
await page.getByText(text, { exact: true }).click();
|
||||
const state = () => page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = app.graph.getNodeById(id);
|
||||
return Object.fromEntries(node.widgets.filter(widget => ['text', 'category', 'tags'].includes(widget.name)).map(widget => [widget.name, widget.value]));
|
||||
}, id);
|
||||
await expect.poll(state).toMatchObject({ text, category: 'Host tests', tags: 'host-test' });
|
||||
for (const renderer of [!testInfo.project.metadata.initialRenderer, testInfo.project.metadata.initialRenderer]) {
|
||||
await page.evaluate(async renderer => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.VueNodes.Enabled', renderer);
|
||||
app.canvas.ds.scale = .8;
|
||||
app.canvas.centerOnNode(app.graph._nodes[0]);
|
||||
}, renderer);
|
||||
await expect.poll(() => page.evaluate(() => LiteGraph.vueNodesMode)).toBe(renderer);
|
||||
await expect(recent).toHaveCount(1);
|
||||
await expect(recent).toBeVisible();
|
||||
expect(await state()).toMatchObject({ text });
|
||||
}
|
||||
const output = await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const prompt = await app.graphToPrompt();
|
||||
const workflow = app.graph.serialize();
|
||||
await app.loadGraphData(workflow);
|
||||
return prompt.output[String(id)].inputs;
|
||||
}, id);
|
||||
expect(output.text).toBe(text);
|
||||
await expect.poll(state).toMatchObject({ text, category: 'Host tests', tags: 'host-test' });
|
||||
await expect(recent).toHaveCount(1);
|
||||
await page.screenshot({ path: testInfo.outputPath(`${type}-reloaded.png`) });
|
||||
const removedStyle = await page.evaluate(async id => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = app.graph.getNodeById(id);
|
||||
const style = node._styleSheet;
|
||||
app.graph.remove(node);
|
||||
return !style.isConnected;
|
||||
}, id);
|
||||
expect(removedStyle).toBe(true);
|
||||
await expect(recent).toHaveCount(0);
|
||||
});
|
||||
}
|
||||
|
||||
test('all node schemas and extension entrypoints load from the real host', async ({ page, request }) => {
|
||||
const definitions = await (await request.get('/object_info')).json();
|
||||
expect(definitions.PromptManager.display_name).toBe('Prompt Manager (Classic)');
|
||||
for (const name of ['PromptManager', 'PromptManagerText', 'PromptSearchList']) expect(definitions[name]?.name).toBe(name);
|
||||
const extensions = await (await request.get('/extensions')).json();
|
||||
expect(extensions.filter(path => path.startsWith('/extensions/promptmanager/'))).toEqual(['/extensions/promptmanager/prompt_manager.js']);
|
||||
const result = await page.evaluate(async () => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = LiteGraph.createNode('PromptSearchList');
|
||||
app.graph.add(node);
|
||||
const workflow = app.graph.serialize();
|
||||
await app.loadGraphData(workflow);
|
||||
return { type: app.graph._nodes[0].type, extension: app.extensions.some(item => item.name === 'PromptManager.UI') };
|
||||
});
|
||||
expect(result).toEqual({ type: 'PromptSearchList', extension: true });
|
||||
});
|
||||
|
||||
for (const type of ['PromptManager', 'PromptManagerText']) {
|
||||
test(`${type} Web UI button opens actual standalone pages`, async ({ page }) => {
|
||||
await page.evaluate(async type => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = LiteGraph.createNode(type);
|
||||
app.graph.add(node);
|
||||
app.canvas.centerOnNode(node);
|
||||
}, type);
|
||||
const popupPromise = page.waitForEvent('popup');
|
||||
await page.getByRole('button', { name: '🌐 Web UI', exact: true }).click();
|
||||
const popup = await popupPromise;
|
||||
await popup.waitForLoadState('domcontentloaded');
|
||||
await expect(popup).toHaveURL(/\/prompt_manager\/(web|admin)/);
|
||||
await expect(popup.locator('body')).toContainText('Prompt');
|
||||
await popup.goto(new URL('/prompt_manager/gallery', page.url()).href);
|
||||
await expect(popup.locator('#addPromptBtn')).toBeVisible();
|
||||
await popup.close();
|
||||
});
|
||||
}
|
||||
|
||||
test('Classic controls follow actual ComfyUI light and dark palettes', async ({ page }, testInfo) => {
|
||||
await page.evaluate(async () => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
const node = LiteGraph.createNode('PromptManager');
|
||||
app.graph.add(node);
|
||||
node.setSize([650, 800]);
|
||||
app.canvas.centerOnNode(node);
|
||||
});
|
||||
for (const palette of ['light', 'dark']) {
|
||||
await page.evaluate(async palette => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
await app.ui.settings.setSettingValue('Comfy.ColorPalette', palette);
|
||||
}, palette);
|
||||
await expect(page.locator('.pm-classic')).toHaveAttribute('data-ca-theme', palette);
|
||||
await expect(page.getByRole('button', { name: '🔍 Search', exact: true })).toHaveCSS('background-color', 'rgb(8, 119, 187)');
|
||||
await page.screenshot({ path: testInfo.outputPath(`classic-${palette}.png`) });
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,50 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const path = require('node:path');
|
||||
const vm = require('node:vm');
|
||||
function fixture() {
|
||||
const panel = { innerHTML: '', isConnected: true };
|
||||
const requests = new Map();
|
||||
const context = vm.createContext({ document: { readyState: 'loading', addEventListener() {},
|
||||
getElementById: id => id === 'metadata-content' ? panel : null },
|
||||
window: {}, console: { error() {} },
|
||||
});
|
||||
const Admin = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/admin.js'), 'utf8') + '\nPromptAdmin;', context, { filename: path.join(__dirname, '../../web/js/admin.js') });
|
||||
const manager = Object.create(Admin.prototype);
|
||||
manager.extractImageMetadata = url => new Promise((resolve, reject) => requests.set(url, { resolve, reject }));
|
||||
manager.updateMetadataSidebar = metadata => { panel.innerHTML = metadata?.imagePath || 'No metadata'; };
|
||||
manager.hideModal = () => {};
|
||||
const load = url => manager.loadMetadataForImage({ getAttribute: () => url });
|
||||
return { manager, panel, requests, load };
|
||||
}
|
||||
test('admin metadata and copy actions stay on the selected image', async () => {
|
||||
const { manager, panel, requests, load } = fixture();
|
||||
const old = load('old.png');
|
||||
const current = load('current.png');
|
||||
requests.get('current.png').resolve({ imagePath: 'current.png' }); await current;
|
||||
requests.get('old.png').resolve({ imagePath: 'old.png' }); await old;
|
||||
assert.equal(panel.innerHTML, 'current.png');
|
||||
assert.equal(manager.currentMetadata.imagePath, 'current.png');
|
||||
});
|
||||
test('navigation clears old action data and late errors cannot replace current metadata', async () => {
|
||||
const { manager, panel, requests, load } = fixture();
|
||||
manager.currentMetadata = { imagePath: 'previous.png' };
|
||||
const old = load('old.png');
|
||||
assert.equal(manager.currentMetadata, null);
|
||||
const current = load('current.png');
|
||||
requests.get('current.png').resolve({ imagePath: 'current.png' }); await current;
|
||||
requests.get('old.png').reject(new Error('old error')); await old;
|
||||
assert.equal(panel.innerHTML, 'current.png');
|
||||
});
|
||||
for (const close of ['close gallery', 'detach sidebar']) {
|
||||
test(`${close} prevents pending admin metadata from restoring copy state`, async () => {
|
||||
const { manager, panel, requests, load } = fixture();
|
||||
const pending = load('old.png');
|
||||
if (close === 'close gallery') manager.closeGallery();
|
||||
else panel.isConnected = false;
|
||||
requests.get('old.png').resolve({ imagePath: 'old.png' }); await pending;
|
||||
assert.equal(manager.currentMetadata, null);
|
||||
assert.notEqual(panel.innerHTML, 'old.png');
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { createApiClient, resolveUrl, fetchUrl, ApiError, latestOnly, STALE } = require("../../web/js/api-client.js");
|
||||
|
||||
test("Prompt Manager URLs follow the page's proxy mount without losing query strings", () => {
|
||||
const url = "/prompt_manager/images/output?limit=100&subfolder=2025%2F04%2F02";
|
||||
assert.equal(resolveUrl(url, "/ComfyBackendDirect/prompt_manager/gallery"),
|
||||
"/ComfyBackendDirect" + url);
|
||||
assert.equal(resolveUrl(url, "/prompt_manager/gallery"), url);
|
||||
assert.equal(resolveUrl(url, "/nested/proxy/prompt_manager/admin"), "/nested/proxy" + url);
|
||||
assert.equal(resolveUrl("/other/api?x=1", "/ComfyBackendDirect/prompt_manager/gallery"), "/other/api?x=1");
|
||||
assert.equal(resolveUrl("/ComfyBackendDirect/prompt_manager/images/serve/a.png", "/ComfyBackendDirect/prompt_manager/gallery"),
|
||||
"/ComfyBackendDirect/prompt_manager/images/serve/a.png");
|
||||
assert.equal(resolveUrl("/prompt_manager/images/serve/a.png?root=1", "/ComfyBackendDirect/prompt_manager/gallery"),
|
||||
"/ComfyBackendDirect/prompt_manager/images/serve/a.png?root=1");
|
||||
});
|
||||
|
||||
test("fetchUrl sends Prompt Manager requests under the current proxy mount", async () => {
|
||||
const previousFetch = globalThis.fetch;
|
||||
const previousLocation = globalThis.location;
|
||||
const calls = [];
|
||||
globalThis.location = { pathname: "/ComfyBackendDirect/prompt_manager/admin" };
|
||||
globalThis.fetch = async (url, init) => {
|
||||
calls.push([url, init]);
|
||||
return { ok: true };
|
||||
};
|
||||
try {
|
||||
await fetchUrl("/prompt_manager/settings?show=all", { method: "GET" });
|
||||
assert.deepEqual(calls, [["/ComfyBackendDirect/prompt_manager/settings?show=all", { method: "GET" }]]);
|
||||
} finally {
|
||||
globalThis.fetch = previousFetch;
|
||||
globalThis.location = previousLocation;
|
||||
}
|
||||
});
|
||||
|
||||
/** A fetch stand-in whose responses the test resolves by hand, in any order. */
|
||||
function fakeFetch() {
|
||||
const calls = [];
|
||||
const impl = (url, init = {}) =>
|
||||
new Promise((resolve, reject) => {
|
||||
const call = { url, init, resolve, reject, aborted: false };
|
||||
if (init.signal) {
|
||||
init.signal.addEventListener("abort", () => {
|
||||
call.aborted = true;
|
||||
reject(new DOMException("The operation was aborted.", "AbortError"));
|
||||
});
|
||||
}
|
||||
calls.push(call);
|
||||
});
|
||||
impl.calls = calls;
|
||||
return impl;
|
||||
}
|
||||
|
||||
function jsonResponse(status, body, { invalid = false } = {}) {
|
||||
return {
|
||||
ok: status >= 200 && status < 300,
|
||||
status,
|
||||
statusText: status === 200 ? "OK" : `Status ${status}`,
|
||||
json: async () => {
|
||||
if (invalid) throw new SyntaxError("Unexpected token");
|
||||
return body;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
test("get returns the parsed body of a successful response", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const pending = api.get("/prompt_manager/recent?page=1");
|
||||
fetch.calls[0].resolve(jsonResponse(200, { success: true, results: [1, 2] }));
|
||||
assert.deepEqual(await pending, { success: true, results: [1, 2] });
|
||||
assert.equal(fetch.calls[0].init.method, "GET");
|
||||
});
|
||||
|
||||
test("get forwards an AbortSignal to fetch", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const controller = new AbortController();
|
||||
const pending = api.get("/x", { signal: controller.signal });
|
||||
assert.equal(fetch.calls[0].init.signal, controller.signal);
|
||||
controller.abort();
|
||||
await assert.rejects(pending, (err) => err.name === "AbortError");
|
||||
});
|
||||
|
||||
test("non-2xx responses throw ApiError carrying the status and server message", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const pending = api.get("/x");
|
||||
fetch.calls[0].resolve(jsonResponse(404, { success: false, error: "Prompt not found" }));
|
||||
await assert.rejects(pending, (err) => {
|
||||
assert.ok(err instanceof ApiError);
|
||||
assert.ok(err instanceof Error);
|
||||
assert.equal(err.status, 404);
|
||||
assert.equal(err.message, "Prompt not found");
|
||||
return true;
|
||||
});
|
||||
});
|
||||
|
||||
test("non-2xx responses without a JSON body fall back to the status text", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const pending = api.get("/x");
|
||||
fetch.calls[0].resolve(jsonResponse(500, null, { invalid: true }));
|
||||
await assert.rejects(pending, (err) => err.status === 500 && err.message === "Status 500");
|
||||
});
|
||||
|
||||
test("a 2xx body with success:false throws ApiError", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const pending = api.post("/x", {});
|
||||
fetch.calls[0].resolve(jsonResponse(200, { success: false, error: "Nope" }));
|
||||
await assert.rejects(pending, (err) => err instanceof ApiError && err.status === 200 && err.message === "Nope");
|
||||
});
|
||||
|
||||
test("a 2xx body that is not JSON throws ApiError", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const pending = api.get("/x");
|
||||
fetch.calls[0].resolve(jsonResponse(200, null, { invalid: true }));
|
||||
await assert.rejects(pending, (err) => err instanceof ApiError && err.status === 200);
|
||||
});
|
||||
|
||||
test("post, put and del send JSON bodies with the right method", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const p1 = api.post("/a", { text: "x" });
|
||||
const p2 = api.put("/b", { rating: 3 });
|
||||
const p3 = api.del("/c");
|
||||
assert.equal(fetch.calls[0].init.method, "POST");
|
||||
assert.equal(fetch.calls[0].init.body, JSON.stringify({ text: "x" }));
|
||||
assert.equal(fetch.calls[0].init.headers["Content-Type"], "application/json");
|
||||
assert.equal(fetch.calls[1].init.method, "PUT");
|
||||
assert.equal(fetch.calls[1].init.body, JSON.stringify({ rating: 3 }));
|
||||
assert.equal(fetch.calls[2].init.method, "DELETE");
|
||||
assert.equal(fetch.calls[2].init.body, undefined);
|
||||
for (const c of fetch.calls) c.resolve(jsonResponse(200, { success: true }));
|
||||
await Promise.all([p1, p2, p3]);
|
||||
});
|
||||
|
||||
test("network failures propagate as ordinary errors", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const pending = api.get("/x");
|
||||
fetch.calls[0].reject(new TypeError("Failed to fetch"));
|
||||
await assert.rejects(pending, (err) => err instanceof TypeError && !(err instanceof ApiError));
|
||||
});
|
||||
|
||||
test("createApiClient without an argument uses the global fetch", async () => {
|
||||
const original = globalThis.fetch;
|
||||
globalThis.fetch = async () => jsonResponse(200, { success: true, via: "global" });
|
||||
try {
|
||||
assert.deepEqual(await createApiClient().get("/x"), { success: true, via: "global" });
|
||||
} finally {
|
||||
globalThis.fetch = original;
|
||||
}
|
||||
});
|
||||
|
||||
test("latestOnly: a newer call aborts the older one and the older result is STALE", async () => {
|
||||
const fetch = fakeFetch();
|
||||
const api = createApiClient(fetch);
|
||||
const latest = latestOnly();
|
||||
|
||||
const first = latest((signal) => api.get("/search?q=a", { signal }));
|
||||
const second = latest((signal) => api.get("/search?q=ab", { signal }));
|
||||
|
||||
assert.equal(fetch.calls[0].aborted, true, "first request aborted");
|
||||
assert.equal(fetch.calls[1].aborted, false);
|
||||
|
||||
fetch.calls[1].resolve(jsonResponse(200, { success: true, q: "ab" }));
|
||||
assert.deepEqual(await second, { success: true, q: "ab" });
|
||||
assert.equal(await first, STALE);
|
||||
});
|
||||
|
||||
test("latestOnly: an older response that arrives after a newer request is STALE even without abort", async () => {
|
||||
const latest = latestOnly();
|
||||
let resolveFirst;
|
||||
const first = latest(() => new Promise((r) => { resolveFirst = r; }));
|
||||
const second = latest(async () => "new");
|
||||
assert.equal(await second, "new");
|
||||
resolveFirst("old");
|
||||
assert.equal(await first, STALE);
|
||||
});
|
||||
|
||||
test("latestOnly: sequential calls each return their own result", async () => {
|
||||
const latest = latestOnly();
|
||||
assert.equal(await latest(async () => 1), 1);
|
||||
assert.equal(await latest(async () => 2), 2);
|
||||
});
|
||||
|
||||
test("latestOnly: errors from the current call propagate, errors from superseded calls do not", async () => {
|
||||
const latest = latestOnly();
|
||||
let rejectFirst;
|
||||
const first = latest(() => new Promise((_, rej) => { rejectFirst = rej; }));
|
||||
const second = latest(async () => { throw new Error("boom"); });
|
||||
await assert.rejects(second, /boom/);
|
||||
rejectFirst(new Error("old failure"));
|
||||
assert.equal(await first, STALE);
|
||||
});
|
||||
@@ -0,0 +1,74 @@
|
||||
// Run with: node --test tests/js/
|
||||
// Each page module is a UMD-style IIFE: CommonJS export under Node, a window
|
||||
// global in the browser. This loads every module the way a browser would and
|
||||
// checks the global name the pages rely on.
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const fs = require("node:fs");
|
||||
const path = require("node:path");
|
||||
const vm = require("node:vm");
|
||||
|
||||
const MODULES = {
|
||||
"html-escape.js": ["HtmlEscape", "escapeHtml", "escapeAttr"],
|
||||
"data-actions.js": ["DataActions"],
|
||||
"api-client.js": ["ApiClient"],
|
||||
"list-state.js": ["ListState"],
|
||||
"sse-stream.js": ["SseStream"],
|
||||
"png-metadata.js": ["PngMetadata"],
|
||||
"view-router.js": ["ViewRouter"],
|
||||
"prompt-list-sort.js": ["PromptListSort"],
|
||||
"comfy-metadata.js": ["ComfyMetadata"],
|
||||
};
|
||||
|
||||
function loadInFakeBrowser(file) {
|
||||
const source = fs.readFileSync(path.join(__dirname, "../../web/js", file), "utf8");
|
||||
const window = {};
|
||||
window.window = window;
|
||||
const context = vm.createContext({
|
||||
window,
|
||||
globalThis: window,
|
||||
console,
|
||||
TextDecoder,
|
||||
TextEncoder,
|
||||
AbortController,
|
||||
URLSearchParams,
|
||||
Symbol,
|
||||
Object,
|
||||
Array,
|
||||
Number,
|
||||
String,
|
||||
Math,
|
||||
Set,
|
||||
Map,
|
||||
JSON,
|
||||
Error,
|
||||
Promise,
|
||||
DataView,
|
||||
Uint8Array,
|
||||
ArrayBuffer,
|
||||
});
|
||||
vm.runInContext(source, context, { filename: file });
|
||||
return window;
|
||||
}
|
||||
|
||||
for (const [file, globals] of Object.entries(MODULES)) {
|
||||
test(`${file} exposes ${globals.join(", ")} on window when module is undefined`, () => {
|
||||
const window = loadInFakeBrowser(file);
|
||||
for (const name of globals) {
|
||||
assert.ok(window[name], `${name} missing`);
|
||||
assert.ok(["object", "function"].includes(typeof window[name]), `${name} is a ${typeof window[name]}`);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
test("the browser escapeHtml global is the same function as HtmlEscape.escapeHtml", () => {
|
||||
const window = loadInFakeBrowser("html-escape.js");
|
||||
assert.equal(window.escapeHtml, window.HtmlEscape.escapeHtml);
|
||||
assert.equal(window.escapeAttr, window.escapeHtml);
|
||||
assert.equal(window.escapeHtml(`<a href="x">`), "<a href="x">");
|
||||
});
|
||||
|
||||
test("the browser ViewRouter global resolves the tags deep link", () => {
|
||||
const window = loadInFakeBrowser("view-router.js");
|
||||
assert.equal(window.ViewRouter.resolveView("#/tags/a"), "tags");
|
||||
});
|
||||
@@ -0,0 +1,363 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { extractGenerationParams } = require("../../web/js/comfy-metadata.js");
|
||||
|
||||
// Same shape as the workflow attached to #75: PromptManager nodes are saved as
|
||||
// CLIPTextEncode with extra inputs, and the checkpoint loader is pysssss'.
|
||||
const issue75Prompt = {
|
||||
1: {
|
||||
class_type: "KSampler",
|
||||
inputs: {
|
||||
seed: 480655185510613, steps: 25, cfg: 5.5,
|
||||
sampler_name: "euler_ancestral", scheduler: "normal", denoise: 0.8,
|
||||
model: ["20", 0], positive: ["39", 0], negative: ["40", 0], latent_image: ["36", 0],
|
||||
},
|
||||
},
|
||||
8: { class_type: "CLIPSetLastLayer", inputs: { stop_at_clip_layer: -2, clip: ["20", 1] } },
|
||||
20: {
|
||||
class_type: "CheckpointLoader|pysssss",
|
||||
inputs: { ckpt_name: "Anime\\Model-XL.safetensors", prompt: "[none]", example: "[none]" },
|
||||
},
|
||||
39: {
|
||||
class_type: "CLIPTextEncode",
|
||||
inputs: {
|
||||
text: "1girl, cat_ears, smile", category: "", tags: "", search_text: "",
|
||||
prepend_text: "masterpiece, best quality, ", append_text: "", clip: ["8", 0],
|
||||
},
|
||||
},
|
||||
40: {
|
||||
class_type: "CLIPTextEncode",
|
||||
inputs: {
|
||||
text: "3d, crowd", category: "", tags: "", search_text: "",
|
||||
prepend_text: "bad quality, worst quality, ", append_text: "", clip: ["8", 0],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
test("#75: traces PromptManager prompts and pysssss checkpoint from the sampler", () => {
|
||||
const r = extractGenerationParams({ prompt: issue75Prompt });
|
||||
assert.equal(r.positivePrompt, "masterpiece, best quality, 1girl, cat_ears, smile");
|
||||
assert.equal(r.negativePrompt, "bad quality, worst quality, 3d, crowd");
|
||||
assert.equal(r.checkpoint, "Anime\\Model-XL.safetensors");
|
||||
assert.deepEqual(
|
||||
[r.seed, r.steps, r.cfgScale, r.sampler],
|
||||
[480655185510613, 25, 5.5, "euler_ancestral"],
|
||||
);
|
||||
});
|
||||
|
||||
test("#75: workflow blob does not override values from the prompt graph", () => {
|
||||
const workflow = { nodes: [{ type: "KSampler", widgets_values: [1, "randomize", 99, 7, "ddim"] }] };
|
||||
const r = extractGenerationParams({ prompt: issue75Prompt, workflow });
|
||||
assert.equal(r.steps, 25);
|
||||
assert.equal(r.cfgScale, 5.5);
|
||||
});
|
||||
|
||||
test("#75: recovers per-node text when SaveImage patch overwrote every PromptManager node", () => {
|
||||
// Pre-fix images: both nodes carry the final positive text in the prompt graph,
|
||||
// but the workflow graph still has each node's real widget values.
|
||||
const positiveFinal = "masterpiece, best quality, 1girl, smile";
|
||||
const prompt = {
|
||||
1: { class_type: "KSampler", inputs: { seed: 1, steps: 20, cfg: 7, sampler_name: "euler", positive: ["39", 0], negative: ["40", 0], model: ["20", 0] } },
|
||||
20: { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: "m.safetensors" } },
|
||||
39: { class_type: "PromptManager", inputs: { text: positiveFinal, prepend_text: "masterpiece, best quality,", append_text: "" } },
|
||||
40: { class_type: "PromptManager", inputs: { text: positiveFinal, prepend_text: "bad quality,", append_text: "" } },
|
||||
};
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 39, type: "PromptManager", widgets_values: ["1girl, smile", "", "", "", "masterpiece, best quality,", ""] },
|
||||
{ id: 40, type: "PromptManager", widgets_values: ["3d, crowd", "", "", "", "bad quality,", ""] },
|
||||
],
|
||||
};
|
||||
const r = extractGenerationParams({ prompt, workflow });
|
||||
assert.equal(r.positivePrompt, "masterpiece, best quality, 1girl, smile");
|
||||
assert.equal(r.negativePrompt, "bad quality, 3d, crowd");
|
||||
});
|
||||
|
||||
test("current metadata: API text wins over a stale or unexpanded workflow widget", () => {
|
||||
// Dynamic prompt expanded at queue time: API holds what ran, workflow holds the template
|
||||
const prompt = {
|
||||
1: { class_type: "KSampler", inputs: { positive: ["39", 0], negative: ["40", 0], model: ["20", 0] } },
|
||||
20: { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: "m.safetensors" } },
|
||||
39: { class_type: "PromptManager", inputs: { text: "a red cat", prepend_text: "", append_text: "" } },
|
||||
40: { class_type: "PromptManager", inputs: { text: "blurry", prepend_text: "", append_text: "" } },
|
||||
};
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 39, type: "PromptManager", widgets_values: ["a {red|blue} cat", "", "", "", "", ""] },
|
||||
{ id: 40, type: "PromptManager", widgets_values: ["blurry", "", "", "", "", ""] },
|
||||
],
|
||||
};
|
||||
const r = extractGenerationParams({ prompt, workflow });
|
||||
assert.equal(r.positivePrompt, "a red cat");
|
||||
assert.equal(r.negativePrompt, "blurry");
|
||||
});
|
||||
|
||||
test("current metadata: linked text is followed, not replaced by the widget's stale value", () => {
|
||||
const prompt = {
|
||||
1: { class_type: "KSampler", inputs: { positive: ["39", 0], negative: ["40", 0], model: ["20", 0] } },
|
||||
20: { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: "m.safetensors" } },
|
||||
39: { class_type: "PromptManager", inputs: { text: ["50", 0], prepend_text: "", append_text: "" } },
|
||||
40: { class_type: "PromptManager", inputs: { text: ["50", 0], prepend_text: "bad,", append_text: "" } },
|
||||
50: { class_type: "PrimitiveString", inputs: { value: "shared text" } },
|
||||
};
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 39, type: "PromptManager", widgets_values: ["old value A", "", "", "", "", ""] },
|
||||
{ id: 40, type: "PromptManager", widgets_values: ["old value B", "", "", "", "bad,", ""] },
|
||||
],
|
||||
};
|
||||
const r = extractGenerationParams({ prompt, workflow });
|
||||
assert.equal(r.positivePrompt, "shared text");
|
||||
});
|
||||
|
||||
test("prepend/append are not doubled when text is already the combined prompt", () => {
|
||||
const prompt = {
|
||||
1: { class_type: "KSampler", inputs: { positive: ["2", 0], negative: ["3", 0], model: ["4", 0] } },
|
||||
2: { class_type: "PromptManager", inputs: { text: "best quality, a cat, 8k", prepend_text: "best quality,", append_text: "8k" } },
|
||||
3: { class_type: "CLIPTextEncode", inputs: { text: "blurry" } },
|
||||
};
|
||||
assert.equal(extractGenerationParams({ prompt }).positivePrompt, "best quality, a cat, 8k");
|
||||
});
|
||||
|
||||
test("linked prepend/append inputs are resolved (no longer baked into text since 3.2.3)", () => {
|
||||
const prompt = {
|
||||
1: { class_type: "KSampler", inputs: { positive: ["2", 0], negative: ["3", 0], model: ["4", 0] } },
|
||||
2: { class_type: "PromptManager", inputs: { text: "a fox", prepend_text: ["5", 0], append_text: ["6", 0] } },
|
||||
3: { class_type: "CLIPTextEncode", inputs: { text: "blurry" } },
|
||||
5: { class_type: "PrimitiveString", inputs: { value: "masterpiece," } },
|
||||
6: { class_type: "PrimitiveString", inputs: { value: "8k" } },
|
||||
};
|
||||
assert.equal(extractGenerationParams({ prompt }).positivePrompt, "masterpiece, a fox 8k");
|
||||
});
|
||||
|
||||
test("negative prompt is identified by link, not by keywords", () => {
|
||||
const prompt = {
|
||||
3: { class_type: "KSampler", inputs: { seed: 1, steps: 20, cfg: 7, sampler_name: "euler", positive: ["6", 0], negative: ["7", 0], model: ["4", 0] } },
|
||||
4: { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: "sd15.safetensors" } },
|
||||
6: { class_type: "CLIPTextEncode", inputs: { text: "embedding:goodstyle, a lighthouse" } },
|
||||
7: { class_type: "CLIPTextEncode", inputs: { text: "blurry" } },
|
||||
};
|
||||
const r = extractGenerationParams({ prompt });
|
||||
assert.equal(r.positivePrompt, "embedding:goodstyle, a lighthouse");
|
||||
assert.equal(r.negativePrompt, "blurry");
|
||||
});
|
||||
|
||||
test("follows LoRA chains, ControlNet pass-through, concatenated text and primitive inputs", () => {
|
||||
const prompt = {
|
||||
1: { class_type: "KSamplerAdvanced", inputs: { noise_seed: 42, steps: ["9", 0], cfg: 6, sampler_name: "dpmpp_2m", positive: ["5", 0], negative: ["5", 1], model: ["3", 0] } },
|
||||
2: { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: "base.safetensors" } },
|
||||
3: { class_type: "LoraLoader", inputs: { lora_name: "style.safetensors", model: ["2", 0], clip: ["2", 1] } },
|
||||
5: { class_type: "ControlNetApplyAdvanced", inputs: { positive: ["6", 0], negative: ["7", 0], strength: 1 } },
|
||||
6: { class_type: "CLIPTextEncode", inputs: { text: ["8", 0] } },
|
||||
7: { class_type: "CLIPTextEncode", inputs: { text: "lowres" } },
|
||||
8: { class_type: "Text Concatenate", inputs: { delimiter: ", ", text_a: "a castle", text_b: "at dusk" } },
|
||||
9: { class_type: "PrimitiveInt", inputs: { value: 30 } },
|
||||
};
|
||||
const r = extractGenerationParams({ prompt });
|
||||
assert.equal(r.positivePrompt, "a castle, at dusk");
|
||||
assert.equal(r.negativePrompt, "lowres");
|
||||
assert.equal(r.checkpoint, "base.safetensors");
|
||||
assert.equal(r.seed, 42);
|
||||
assert.equal(r.steps, 30);
|
||||
});
|
||||
|
||||
test("custom sampler graphs (Flux: SamplerCustomAdvanced + CFGGuider) resolve every field", () => {
|
||||
const prompt = {
|
||||
85: { class_type: "UNETLoader", inputs: { unet_name: "flux2-klein.safetensors" } },
|
||||
135: { class_type: "Lora Loader (LoraManager)", inputs: { model: ["85", 0], clip: ["90", 0] } },
|
||||
149: { class_type: "PromptManager", inputs: { text: "a red fox", prepend_text: ["133", 0], append_text: "" } },
|
||||
152: { class_type: "SeedHistory", inputs: { seed: 1725922648 } },
|
||||
// ControlNet-style node with positive/negative but no model: must not be taken as the sampler
|
||||
200: { class_type: "ControlNetApplyAdvanced", inputs: { positive: ["149", 0], negative: ["274", 0], strength: 1 } },
|
||||
265: { class_type: "SamplerCustomAdvanced", inputs: { noise: ["267", 0], guider: ["268", 0], sampler: ["269", 0], sigmas: ["266", 0] } },
|
||||
266: { class_type: "Flux2Scheduler", inputs: { steps: 50, width: ["153", 1] } },
|
||||
267: { class_type: "RandomNoise", inputs: { noise_seed: ["152", 0] } },
|
||||
268: { class_type: "CFGGuider", inputs: { cfg: 5.0, model: ["135", 0], positive: ["149", 0], negative: ["274", 0] } },
|
||||
269: { class_type: "KSamplerSelect", inputs: { sampler_name: "uni_pc" } },
|
||||
274: { class_type: "CLIPTextEncode", inputs: { text: "watermark" } },
|
||||
};
|
||||
const r = extractGenerationParams({ prompt });
|
||||
assert.deepEqual(
|
||||
[r.positivePrompt, r.negativePrompt, r.checkpoint, r.cfgScale, r.seed, r.steps, r.sampler],
|
||||
["a red fox", "watermark", "flux2-klein.safetensors", 5.0, 1725922648, 50, "uni_pc"],
|
||||
);
|
||||
});
|
||||
|
||||
test("multi-output settings nodes (SamplerCombo) map each field to its own value", () => {
|
||||
const prompt = {
|
||||
5: {
|
||||
class_type: "KSampler",
|
||||
inputs: {
|
||||
seed: ["162", 0], steps: ["165", 2], cfg: ["165", 3], sampler_name: ["165", 0],
|
||||
positive: ["6", 0], negative: ["7", 0], model: ["4", 0],
|
||||
},
|
||||
},
|
||||
4: { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: "m.safetensors" } },
|
||||
6: { class_type: "CLIPTextEncode", inputs: { text: "a cat" } },
|
||||
7: { class_type: "CLIPTextEncode", inputs: { text: "blurry" } },
|
||||
162: { class_type: "SeedHistory", inputs: { seed: 1495646836, seed_history_ui: "" } },
|
||||
165: { class_type: "SamplerCombo", inputs: { sampler_name: "euler_ancestral", scheduler: "normal", steps: 50, cfg: 7 } },
|
||||
};
|
||||
const r = extractGenerationParams({ prompt });
|
||||
assert.deepEqual([r.seed, r.steps, r.cfgScale, r.sampler], [1495646836, 50, 7, "euler_ancestral"]);
|
||||
});
|
||||
|
||||
test("workflow-only images account for the hidden control_after_generate widget", () => {
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 1, type: "CheckpointLoader|pysssss", widgets_values: ["model.safetensors", "[none]"] },
|
||||
{ id: 2, type: "KSampler", widgets_values: [123, "randomize", 25, 5.5, "euler", "normal", 1] },
|
||||
],
|
||||
};
|
||||
const r = extractGenerationParams({ workflow });
|
||||
assert.deepEqual(
|
||||
[r.checkpoint, r.seed, r.steps, r.cfgScale, r.sampler],
|
||||
["model.safetensors", 123, 25, 5.5, "euler"],
|
||||
);
|
||||
});
|
||||
|
||||
test("cyclic, fan-out graphs from crafted PNGs resolve quickly", () => {
|
||||
// Every node links to itself and its neighbour through several inputs and has
|
||||
// no scalars: without a visited set this explores ~inputs^MAX_DEPTH paths.
|
||||
const loop = (id, next) => ({
|
||||
class_type: "Evil",
|
||||
inputs: {
|
||||
text_a: [id, 0], text_b: [next, 0], text_c: [id, 0], text_d: [next, 0],
|
||||
value_a: [id, 0], value_b: [next, 0], model: [next, 0],
|
||||
},
|
||||
});
|
||||
const prompt = {
|
||||
1: { class_type: "KSampler", inputs: { seed: ["2", 0], steps: ["3", 0], positive: ["2", 0], negative: ["3", 0], model: ["2", 0] } },
|
||||
2: loop("2", "3"),
|
||||
3: loop("3", "2"),
|
||||
};
|
||||
const started = Date.now();
|
||||
const r = extractGenerationParams({ prompt });
|
||||
assert.ok(Date.now() - started < 500, `took ${Date.now() - started}ms`);
|
||||
assert.equal(r.positivePrompt, "No prompt found");
|
||||
assert.equal(r.seed, "Unknown");
|
||||
});
|
||||
|
||||
test("chained encoders whose prepend and append both link onward stay linear", () => {
|
||||
const prompt = { 1: { class_type: "KSampler", inputs: { positive: ["n0", 0], negative: ["n0", 0], model: ["n0", 0] } } };
|
||||
for (let i = 0; i < 60; i++) {
|
||||
prompt[`n${i}`] = {
|
||||
class_type: "PromptManager",
|
||||
inputs: { text: `t${i}`, prepend_text: [`n${i + 1}`, 0], append_text: [`n${i + 1}`, 0] },
|
||||
};
|
||||
}
|
||||
const started = Date.now();
|
||||
extractGenerationParams({ prompt });
|
||||
assert.ok(Date.now() - started < 500, `took ${Date.now() - started}ms`);
|
||||
});
|
||||
|
||||
// UI-format link: [id, from_node, from_slot, to_node, to_slot, type]
|
||||
test("workflow-only: traces prompts, model and KSamplerAdvanced settings through Set/Get and Reroute", () => {
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 1, type: "UNETLoader", widgets_values: ["flux.safetensors", "default"], outputs: [{ links: [1] }] },
|
||||
{ id: 2, type: "SetNode", widgets_values: ["model"], inputs: [{ name: "MODEL", link: 1 }] },
|
||||
{ id: 3, type: "GetNode", widgets_values: ["model"], outputs: [{ links: [2] }] },
|
||||
{ id: 4, type: "PromptManager", widgets_values: ["a castle", "", "", "", "masterpiece,", ""], inputs: [], outputs: [{ links: [3] }] },
|
||||
{ id: 5, type: "Reroute", inputs: [{ name: "", link: 3 }], outputs: [{ links: [4] }] },
|
||||
{ id: 6, type: "CLIPTextEncode", widgets_values: ["lowres"], outputs: [{ links: [5] }] },
|
||||
{
|
||||
id: 7, type: "KSamplerAdvanced",
|
||||
inputs: [{ name: "model", link: 2 }, { name: "positive", link: 4 }, { name: "negative", link: 5 }],
|
||||
widgets_values: ["enable", 77, "randomize", 30, 6.5, "dpmpp_2m", "karras", 0, 10000, "disable"],
|
||||
},
|
||||
],
|
||||
links: [
|
||||
[1, 1, 0, 2, 0, "MODEL"], [2, 3, 0, 7, 0, "MODEL"], [3, 4, 0, 5, 0, "CONDITIONING"],
|
||||
[4, 5, 0, 7, 1, "CONDITIONING"], [5, 6, 0, 7, 2, "CONDITIONING"],
|
||||
],
|
||||
};
|
||||
const r = extractGenerationParams({ workflow });
|
||||
assert.deepEqual(
|
||||
[r.positivePrompt, r.negativePrompt, r.checkpoint, r.seed, r.steps, r.cfgScale, r.sampler],
|
||||
["masterpiece, a castle", "lowres", "flux.safetensors", 77, 30, 6.5, "dpmpp_2m"],
|
||||
);
|
||||
});
|
||||
|
||||
test("workflow-only: widgets_values_named and widget-backed inputs take precedence over positions", () => {
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 1, type: "CheckpointLoaderSimple", widgets_values: ["m.safetensors"], outputs: [{ links: [1] }] },
|
||||
{ id: 2, type: "CLIPTextEncode", widgets_values_named: { text: "named text" }, widgets_values: ["ignored"], outputs: [{ links: [2] }] },
|
||||
{ id: 3, type: "CLIPTextEncode", widgets_values: ["neg"], outputs: [{ links: [3] }] },
|
||||
{
|
||||
id: 4, type: "KSampler",
|
||||
inputs: [
|
||||
{ name: "model", link: 1 }, { name: "positive", link: 2 }, { name: "negative", link: 3 },
|
||||
{ name: "seed", widget: { name: "seed" }, link: null },
|
||||
{ name: "steps", widget: { name: "steps" }, link: null },
|
||||
{ name: "cfg", widget: { name: "cfg" }, link: null },
|
||||
{ name: "sampler_name", widget: { name: "sampler_name" }, link: null },
|
||||
],
|
||||
widgets_values: [5, "fixed", 12, 4, "euler", "normal", 1],
|
||||
},
|
||||
],
|
||||
links: [[1, 1, 0, 4, 0, "MODEL"], [2, 2, 0, 4, 1, "CONDITIONING"], [3, 3, 0, 4, 2, "CONDITIONING"]],
|
||||
};
|
||||
const r = extractGenerationParams({ workflow });
|
||||
assert.deepEqual([r.positivePrompt, r.seed, r.steps, r.cfgScale, r.sampler], ["named text", 5, 12, 4, "euler"]);
|
||||
});
|
||||
|
||||
test("workflow-only: bypassed nodes pass through, muted nodes are ignored", () => {
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 1, type: "CheckpointLoaderSimple", widgets_values: ["m.safetensors"], outputs: [{ links: [1] }] },
|
||||
{ id: 2, type: "LoraLoader", mode: 4, inputs: [{ name: "model", type: "MODEL", link: 1 }], outputs: [{ type: "MODEL", links: [2] }] },
|
||||
{ id: 3, type: "CLIPTextEncode", widgets_values: ["pos"], outputs: [{ links: [3] }] },
|
||||
{ id: 4, type: "CLIPTextEncode", widgets_values: ["neg"], outputs: [{ links: [4] }] },
|
||||
{ id: 5, type: "KSampler", mode: 2, inputs: [{ name: "model", link: 2 }, { name: "positive", link: 3 }, { name: "negative", link: 4 }], widgets_values: [9, "fixed", 99, 9, "ddim"] },
|
||||
{ id: 6, type: "KSampler", inputs: [{ name: "model", link: 2 }, { name: "positive", link: 3 }, { name: "negative", link: 4 }], widgets_values: [1, "fixed", 20, 7, "euler"] },
|
||||
],
|
||||
links: [[1, 1, 0, 2, 0, "MODEL"], [2, 2, 0, 6, 0, "MODEL"], [3, 3, 0, 6, 1, "CONDITIONING"], [4, 4, 0, 6, 2, "CONDITIONING"]],
|
||||
};
|
||||
const r = extractGenerationParams({ workflow });
|
||||
assert.deepEqual([r.checkpoint, r.positivePrompt, r.steps], ["m.safetensors", "pos", 20]);
|
||||
});
|
||||
|
||||
test("missing or empty metadata yields placeholders", () => {
|
||||
const r = extractGenerationParams({});
|
||||
assert.equal(r.positivePrompt, "No prompt found");
|
||||
assert.equal(r.negativePrompt, "No negative prompt found");
|
||||
assert.equal(r.checkpoint, "Unknown");
|
||||
assert.equal(extractGenerationParams(null).seed, "Unknown");
|
||||
});
|
||||
|
||||
|
||||
test("Focused API metadata composes linked positive text and preserves negative role", () => {
|
||||
const prompt = {
|
||||
1: { class_type: "KSampler", inputs: { positive: ["2", 0], negative: ["3", 0], model: ["4", 0] } },
|
||||
2: { class_type: "PromptManagerFocused", inputs: { text: "a fox", prepend_text: ["5", 0], append_text: "8k" } },
|
||||
3: { class_type: "PromptManagerFocused", inputs: { text: "blurry", prepend_text: "avoid" } },
|
||||
4: { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: "model.safetensors" } },
|
||||
5: { class_type: "PrimitiveString", inputs: { value: "lora trigger" } },
|
||||
};
|
||||
const result = extractGenerationParams({ prompt });
|
||||
assert.equal(result.positivePrompt, "lora trigger a fox 8k");
|
||||
assert.equal(result.negativePrompt, "avoid blurry");
|
||||
});
|
||||
|
||||
test("Focused workflow-only metadata recovers serialized fields and connected composition", () => {
|
||||
const workflow = {
|
||||
nodes: [
|
||||
{ id: 1, type: "KSampler", inputs: [{ name: "positive", link: 1 }, { name: "negative", link: 2 }, { name: "model", link: 3 }] },
|
||||
{ id: 2, type: "PromptManagerFocused", widgets_values: ["a fox", "animals", "fox", "ignored search", "stale widget", "8k"], inputs: [{ name: "prepend_text", link: 4 }] },
|
||||
{ id: 3, type: "PromptManagerFocused", widgets_values: ["blurry", "negative", "bad", "ignored search", "avoid", "watermark"] },
|
||||
{ id: 4, type: "UNETLoader", widgets_values: ["model.safetensors", "default"] },
|
||||
{ id: 5, type: "CLIPTextEncode", widgets_values: ["lora trigger"] },
|
||||
],
|
||||
links: [
|
||||
[1, 2, 0, 1, 0, "CONDITIONING"], [2, 3, 0, 1, 1, "CONDITIONING"],
|
||||
[3, 4, 0, 1, 2, "MODEL"], [4, 5, 0, 2, 0, "STRING"],
|
||||
],
|
||||
};
|
||||
const result = extractGenerationParams({ workflow });
|
||||
assert.equal(result.positivePrompt, "lora trigger a fox 8k");
|
||||
assert.equal(result.negativePrompt, "avoid blurry watermark");
|
||||
assert.equal(result.checkpoint, "model.safetensors");
|
||||
});
|
||||
@@ -0,0 +1,49 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { readAction, dispatch } = require("../../web/js/data-actions.js");
|
||||
|
||||
test("readAction parses promptId and index as integers", () => {
|
||||
const parsed = readAction({ action: "open-film", promptId: "42", index: "3" });
|
||||
assert.equal(parsed.action, "open-film");
|
||||
assert.equal(parsed.promptId, 42);
|
||||
assert.equal(parsed.index, 3);
|
||||
});
|
||||
|
||||
test("readAction turns missing or non-numeric ids into null", () => {
|
||||
assert.equal(readAction({ action: "edit" }).promptId, null);
|
||||
assert.equal(readAction({ action: "edit", promptId: "abc" }).promptId, null);
|
||||
assert.equal(readAction({ action: "edit", index: "" }).index, null);
|
||||
});
|
||||
|
||||
test("readAction passes string fields through untouched", () => {
|
||||
const parsed = readAction({ action: "remove-tag", tag: `a "quoted" tag`, url: "/x?y=1&z=2", copyText: "p" });
|
||||
assert.equal(parsed.tag, `a "quoted" tag`);
|
||||
assert.equal(parsed.url, "/x?y=1&z=2");
|
||||
assert.equal(parsed.copyText, "p");
|
||||
});
|
||||
|
||||
test("readAction tolerates a missing dataset", () => {
|
||||
assert.equal(readAction(undefined).action, "");
|
||||
assert.equal(readAction(null).promptId, null);
|
||||
});
|
||||
|
||||
test("dispatch calls the matching handler with the parsed action and extras", () => {
|
||||
const calls = [];
|
||||
const handled = dispatch(
|
||||
{ action: "edit", promptId: "7" },
|
||||
{ edit: (a, el) => calls.push([a.promptId, el]) },
|
||||
"element",
|
||||
);
|
||||
assert.equal(handled, true);
|
||||
assert.deepEqual(calls, [[7, "element"]]);
|
||||
});
|
||||
|
||||
test("dispatch returns false and calls nothing for an unknown or empty action", () => {
|
||||
let called = 0;
|
||||
const handlers = { edit: () => called++ };
|
||||
assert.equal(dispatch({ action: "other" }, handlers), false);
|
||||
assert.equal(dispatch({}, handlers), false);
|
||||
assert.equal(dispatch({ action: "constructor" }, handlers), false, "no prototype lookups");
|
||||
assert.equal(called, 0);
|
||||
});
|
||||
@@ -0,0 +1,76 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const vm = require('node:vm');
|
||||
const path = require('node:path');
|
||||
|
||||
function deferred() {
|
||||
let resolve, reject;
|
||||
const promise = new Promise((yes, no) => { resolve = yes; reject = no; });
|
||||
return { promise, resolve, reject };
|
||||
}
|
||||
|
||||
function fixture() {
|
||||
const context = vm.createContext({ document: { addEventListener() {}, getElementById() {} },
|
||||
window: {}, console: { error() {} }, escapeHtml: value => value });
|
||||
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
|
||||
const manager = Object.create(Manager.prototype);
|
||||
const requests = new Map();
|
||||
manager.extractImageMetadata = url => {
|
||||
const pending = deferred();
|
||||
requests.set(url, pending);
|
||||
return pending.promise;
|
||||
};
|
||||
manager.displayMetadata = (metadata, panel) => {
|
||||
panel.innerHTML = metadata?.positivePrompt || 'No metadata';
|
||||
manager.currentMetadata = metadata;
|
||||
};
|
||||
manager.stopViewerSlideshow = () => {};
|
||||
return { manager, requests, panel: { innerHTML: '', isConnected: true } };
|
||||
}
|
||||
|
||||
test('a late response cannot overwrite the selected image or copy state', async () => {
|
||||
const { manager, requests, panel } = fixture();
|
||||
const first = manager.loadImageMetadata('old.png', panel);
|
||||
const second = manager.loadImageMetadata('new.png', panel);
|
||||
requests.get('new.png').resolve({ positivePrompt: 'new prompt' });
|
||||
await second;
|
||||
requests.get('old.png').resolve({ positivePrompt: 'old prompt' });
|
||||
await first;
|
||||
assert.equal(panel.innerHTML, 'new prompt');
|
||||
assert.equal(manager.currentMetadata.positivePrompt, 'new prompt');
|
||||
});
|
||||
|
||||
test('navigation clears old action data while the new request is pending', async () => {
|
||||
const { manager, requests, panel } = fixture();
|
||||
manager.currentMetadata = { positivePrompt: 'old prompt' };
|
||||
const pending = manager.loadImageMetadata('new.png', panel);
|
||||
assert.equal(manager.currentMetadata, null);
|
||||
requests.get('new.png').resolve(null);
|
||||
await pending;
|
||||
assert.equal(manager.currentMetadata, null);
|
||||
});
|
||||
|
||||
test('a stale error cannot replace newer metadata', async () => {
|
||||
const { manager, requests, panel } = fixture();
|
||||
const first = manager.loadImageMetadata('old.png', panel);
|
||||
const second = manager.loadImageMetadata('new.png', panel);
|
||||
requests.get('new.png').resolve({ positivePrompt: 'new prompt' });
|
||||
await second;
|
||||
requests.get('old.png').reject(new Error('old request failed'));
|
||||
await first;
|
||||
assert.equal(panel.innerHTML, 'new prompt');
|
||||
});
|
||||
|
||||
for (const close of ['destroy viewer', 'remove modal']) {
|
||||
test(`closing via ${close} prevents pending metadata from restoring action data`, async () => {
|
||||
const { manager, requests, panel } = fixture();
|
||||
const pending = manager.loadImageMetadata('old.png', panel);
|
||||
if (close === 'destroy viewer') manager.destroyViewer();
|
||||
else panel.isConnected = false;
|
||||
requests.get('old.png').resolve({ positivePrompt: 'closed prompt' });
|
||||
await pending;
|
||||
assert.equal(manager.currentMetadata, null);
|
||||
assert.notEqual(panel.innerHTML, 'closed prompt');
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const vm = require('node:vm');
|
||||
const path = require('node:path');
|
||||
const { escapeHtml } = require('../../web/js/html-escape.js');
|
||||
|
||||
function fixture(count = 2000) {
|
||||
const elements = new Map();
|
||||
function element() {
|
||||
return {
|
||||
dataset: {}, children: [], listeners: {}, classList: { add() {}, remove() {} },
|
||||
set innerHTML(html) { this.html = html; this.children = [{ html }]; },
|
||||
get innerHTML() { return this.html || ''; },
|
||||
insertAdjacentHTML(position, html) { this.children.push({ html }); },
|
||||
addEventListener(type, handler) { (this.listeners[type] ||= []).push(handler); },
|
||||
querySelectorAll() { return []; },
|
||||
remove() {},
|
||||
};
|
||||
}
|
||||
const document = {
|
||||
addEventListener() {},
|
||||
getElementById(id) { if (!elements.has(id)) elements.set(id, element()); return elements.get(id); },
|
||||
createElement: element,
|
||||
};
|
||||
const instances = [];
|
||||
const timers = new Map();
|
||||
let timerId = 0;
|
||||
class Viewer {
|
||||
constructor(holder, options) {
|
||||
this.element = holder; this.options = { fullscreen: true, ...options };
|
||||
this.index = options.initialViewIndex; this.isShown = true;
|
||||
this.length = (holder.innerHTML.match(/<img /g) || []).length;
|
||||
instances.push(this);
|
||||
}
|
||||
requestFullscreen() { this.fullscreenRequested = true; }
|
||||
exitFullscreen() { this.fullscreenExited = true; }
|
||||
show() { this.shown = true; }
|
||||
view(index) { this.index = index; }
|
||||
clearEnforceFocus() { this.focusCleared = true; }
|
||||
destroy() { this.destroyed = true; }
|
||||
}
|
||||
const context = vm.createContext({ document, window: {}, Viewer, escapeHtml,
|
||||
console, setTimeout(callback) { timers.set(++timerId, callback); return timerId; },
|
||||
clearTimeout(id) { timers.delete(id); },
|
||||
ApiClient: { resolveUrl: value => value, STALE: Symbol() },
|
||||
ListState: { nextPageState: (state, next) => ({ ...state, ...next }) },
|
||||
});
|
||||
const source = fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8');
|
||||
const Manager = vm.runInContext(source + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
|
||||
const manager = Object.create(Manager.prototype);
|
||||
Object.assign(manager, { images: Array.from({ length: count }, (_, i) => ({
|
||||
filename: `${i}.png`, url: `/original/${i}.png`, thumbnail_url: `/thumb/${i}.jpg`,
|
||||
relative_path: `${i}.png`, size: 100, modified_time: 1,
|
||||
})), viewMode: 'grid', viewer: null, currentPage: 1, limit: 100, total: count,
|
||||
getSettings: () => ({ showImageInfo: true }), });
|
||||
return { manager, document, instances, timers };
|
||||
}
|
||||
|
||||
for (const mode of ['grid', 'list']) {
|
||||
test(`${mode} append retains existing cards and adds only the new page`, () => {
|
||||
const { manager, document, timers } = fixture(100);
|
||||
manager.viewMode = mode;
|
||||
manager.renderGallery();
|
||||
const container = document.getElementById(mode === 'grid' ? 'galleryGrid' : 'galleryList');
|
||||
assert.equal(timers.size, 0);
|
||||
const first = container.children[0];
|
||||
manager.images.push({ ...manager.images[0], filename: 'new.png' });
|
||||
manager.renderGallery({ from: 100 });
|
||||
assert.equal(container.children[0], first);
|
||||
assert.equal(container.children.length, 2);
|
||||
assert.match(container.children[1].html, /new.png/);
|
||||
assert.doesNotMatch(container.children[1].html, /alt="0.png"/);
|
||||
assert.match(container.children[1].html, /data-image-index="100"/);
|
||||
});
|
||||
}
|
||||
|
||||
test('loadImages passes the previous length for append, and zero for replacement', async () => {
|
||||
const { manager } = fixture(200);
|
||||
const offsets = [];
|
||||
Object.assign(manager, { imagesRequest: async () => ({ total: 500, images: [manager.images[0]] }),
|
||||
renderGallery: options => offsets.push(options.from), updateStats() {},
|
||||
updatePagination() {}, maybeLoadMore() {}, showLoading() {}, });
|
||||
await manager.loadImages(3, { append: true });
|
||||
await manager.loadImages(1);
|
||||
assert.deepEqual(offsets, [200, 0]);
|
||||
assert.equal(manager.images.length, 1);
|
||||
});
|
||||
|
||||
test('viewer opening is bounded and selects the clicked original', () => {
|
||||
const { manager, instances } = fixture();
|
||||
manager.openViewer(1500);
|
||||
const viewer = instances[0];
|
||||
assert.ok(viewer.element.innerHTML.match(/<img /g).length <= 100);
|
||||
assert.match(viewer.element.innerHTML, /data-original="\/original\/1500.png"/);
|
||||
const urls = [...viewer.element.innerHTML.matchAll(/data-original="([^"]+)"/g)].map(m => m[1]);
|
||||
assert.equal(urls[viewer.options.initialViewIndex], '/original/1500.png');
|
||||
assert.equal(viewer.shown, true);
|
||||
});
|
||||
|
||||
test('previous/next cross window edges and wrap only at gallery edges', () => {
|
||||
const { manager, instances } = fixture();
|
||||
manager.openViewer(1500);
|
||||
const first = manager.viewer;
|
||||
first.index = 99;
|
||||
const lastUrl = [...first.element.innerHTML.matchAll(/data-original="\/original\/(\d+).png"/g)].at(-1);
|
||||
first.next(true);
|
||||
const next = manager.viewer;
|
||||
const urls = [...next.element.innerHTML.matchAll(/data-original="\/original\/(\d+).png"/g)];
|
||||
assert.equal(Number(urls[next.index][1]), Number(lastUrl[1]) + 1);
|
||||
manager.openViewer(0);
|
||||
manager.viewer.prev(true);
|
||||
const wrapped = manager.viewer;
|
||||
const wrappedUrls = [...wrapped.element.innerHTML.matchAll(/data-original="\/original\/(\d+).png"/g)];
|
||||
assert.equal(Number(wrappedUrls[wrapped.index][1]), 1999);
|
||||
assert.ok(instances.length >= 3);
|
||||
});
|
||||
|
||||
test('delegated card clicks bind once, open appended indices, and preserve video dispatch', () => {
|
||||
const { manager, document } = fixture();
|
||||
const opened = [];
|
||||
manager.openViewer = index => opened.push(index);
|
||||
manager.showVideoModal = (...args) => opened.push(args);
|
||||
manager.addFallbackClickHandlers();
|
||||
manager.addFallbackClickHandlers();
|
||||
const grid = document.getElementById('galleryGrid');
|
||||
assert.equal(grid.listeners.click.length, 1);
|
||||
const item = { dataset: { imageIndex: '1500' }, querySelector: () => ({ dataset: {
|
||||
original: '/original/1500.png', caption: 'caption', isVideo: 'false',
|
||||
} }) };
|
||||
const click = () => grid.listeners.click[0]({ target: { closest: () => item }, preventDefault() {}, stopPropagation() {} });
|
||||
click();
|
||||
item.querySelector = () => ({ dataset: { original: '/video.mp4', caption: 'video', isVideo: 'true' } });
|
||||
click();
|
||||
assert.deepEqual(opened, [1500, ['/video.mp4', 'video', 1500]]);
|
||||
});
|
||||
|
||||
test('navigation inside the window reuses the viewer and honors non-looping edges', () => {
|
||||
const { manager, instances } = fixture();
|
||||
manager.openViewer(1500);
|
||||
const viewer = manager.viewer;
|
||||
viewer.next(true);
|
||||
assert.equal(manager.viewer, viewer);
|
||||
assert.equal(viewer.index, 51);
|
||||
viewer.prev(true);
|
||||
assert.equal(viewer.index, 50);
|
||||
manager.openViewer(0);
|
||||
const edge = manager.viewer;
|
||||
edge.prev(false);
|
||||
assert.equal(manager.viewer, edge);
|
||||
assert.equal(edge.index, 0);
|
||||
assert.equal(instances.length, 2);
|
||||
});
|
||||
|
||||
test('append leaves an open viewer intact; a full render destroys it', () => {
|
||||
const { manager } = fixture();
|
||||
manager.openViewer(1500);
|
||||
const viewer = manager.viewer;
|
||||
manager.renderGallery({ from: 1900 });
|
||||
assert.equal(manager.viewer, viewer);
|
||||
manager.renderGallery();
|
||||
assert.equal(viewer.destroyed, true);
|
||||
assert.equal(manager.viewer, null);
|
||||
});
|
||||
|
||||
test('each viewed image updates sidebar metadata without waiting for DOM mutations', () => {
|
||||
const { manager, document } = fixture();
|
||||
const urls = [];
|
||||
document.querySelector = () => ({});
|
||||
const sidebar = document.getElementById('metadata-sidebar');
|
||||
sidebar.querySelector = () => ({});
|
||||
manager.loadImageMetadata = url => urls.push(url);
|
||||
manager.openViewer(1500);
|
||||
manager.viewer.options.viewed({ detail: { originalImage: { dataset: { original: '/original/1500.png' } } } });
|
||||
manager.viewer.options.viewed({ detail: { originalImage: { dataset: { original: '/original/1501.png' } } } });
|
||||
assert.deepEqual(urls, ['/original/1500.png', '/original/1501.png']);
|
||||
});
|
||||
|
||||
test('slideshow advances across windows without preloading every original and stops on hide', () => {
|
||||
const { manager, timers } = fixture();
|
||||
manager.openViewer(1500);
|
||||
let prevented = false;
|
||||
manager.viewer.options.play({ preventDefault() { prevented = true; } });
|
||||
assert.equal(prevented, true);
|
||||
assert.equal(manager.viewer.fullscreenRequested, true);
|
||||
assert.equal(timers.size, 1);
|
||||
manager.viewer.index = 99;
|
||||
const tick = timers.values().next().value;
|
||||
timers.clear();
|
||||
tick();
|
||||
const viewer = manager.viewer;
|
||||
const urls = [...viewer.element.innerHTML.matchAll(/data-original="([^"]+)"/g)];
|
||||
assert.equal(urls[viewer.index][1], '/original/1550.png');
|
||||
assert.equal(timers.size, 1);
|
||||
viewer.options.hide();
|
||||
assert.equal(viewer.fullscreenExited, true);
|
||||
assert.equal(timers.size, 0);
|
||||
});
|
||||
|
||||
test('destroying a viewer removes its focus trap and slideshow timer', () => {
|
||||
const { manager, timers } = fixture();
|
||||
manager.openViewer(1500);
|
||||
const viewer = manager.viewer;
|
||||
viewer.options.play({ preventDefault() {} });
|
||||
manager.destroyViewer();
|
||||
assert.equal(viewer.focusCleared, true);
|
||||
assert.equal(viewer.destroyed, true);
|
||||
assert.equal(timers.size, 0);
|
||||
});
|
||||
@@ -0,0 +1,88 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const path = require('node:path');
|
||||
const vm = require('node:vm');
|
||||
const ImageHelpers = require('../../web/js/image-helpers.js');
|
||||
function deferred() {
|
||||
let resolve;
|
||||
const promise = new Promise(yes => { resolve = yes; });
|
||||
return { promise, resolve };
|
||||
}
|
||||
function fixture() {
|
||||
const elements = new Map();
|
||||
const requests = [];
|
||||
const document = { addEventListener() {}, getElementById(id) {
|
||||
if (!elements.has(id)) elements.set(id, {});
|
||||
return elements.get(id);
|
||||
} };
|
||||
const context = vm.createContext({ document, window: {}, ImageHelpers,
|
||||
escapeHtml: value => value, console: { error() {} },
|
||||
ApiClient: { resolveUrl: value => value, fetchUrl: (url, options) => {
|
||||
const request = deferred();
|
||||
requests.push({ url, body: JSON.parse(options.body), ...request });
|
||||
return request.promise;
|
||||
} },
|
||||
});
|
||||
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
|
||||
const manager = Object.create(Manager.prototype);
|
||||
const notifications = [];
|
||||
Object.assign(manager, { autoTagState: { reviewImages: [
|
||||
{ id: 'digest', path: 'same.png', root_index: 1, url: '/same.png' },
|
||||
{ id: 'other', path: 'next.png', root_index: 0, url: '/next.png' },
|
||||
], reviewIndex: 0, currentTags: [], modelType: 'gguf' },
|
||||
renderReviewTags() {}, showNotification: (...args) => notifications.push(args), hideModal() {}, loadImages() {},
|
||||
});
|
||||
const reply = (index, data) => requests[index].resolve({ json: async () => data });
|
||||
return { manager, requests, reply, notifications };
|
||||
}
|
||||
test('gallery review retains root identity and applies tags with the returned prompt ID', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const load = manager.loadNextReviewImage();
|
||||
assert.equal(requests[0].body.root_index, 1);
|
||||
reply(0, { success: true, prompt_id: 42, tags: ['smile'] });
|
||||
await load;
|
||||
manager.loadNextReviewImage = () => {};
|
||||
const apply = manager.applyReviewTags();
|
||||
assert.deepEqual(requests[1].body, { prompt_id: 42, tags: ['smile'] });
|
||||
reply(1, { success: true });
|
||||
await apply;
|
||||
});
|
||||
test('unlinked images do not submit an invalid apply request or advance', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const load = manager.loadNextReviewImage();
|
||||
reply(0, { success: true, prompt_id: null, tags: ['smile'] });
|
||||
await load;
|
||||
const apply = manager.applyReviewTags();
|
||||
if (requests[1]) reply(1, { success: false, error: "Unexpected apply" });
|
||||
await apply;
|
||||
assert.equal(requests.length, 1);
|
||||
assert.equal(manager.autoTagState.reviewIndex, 0);
|
||||
});
|
||||
test('a failed apply keeps the current image available for retry', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const load = manager.loadNextReviewImage();
|
||||
reply(0, { success: true, prompt_id: 42, tags: ['smile'] });
|
||||
await load;
|
||||
const apply = manager.applyReviewTags();
|
||||
reply(1, { success: false, error: 'Try again' });
|
||||
await apply;
|
||||
assert.equal(manager.autoTagState.reviewIndex, 0);
|
||||
assert.equal(requests.length, 2);
|
||||
});
|
||||
test('late generation responses cannot restore tags from a skipped or closed image', async () => {
|
||||
const { manager, reply } = fixture();
|
||||
const first = manager.loadNextReviewImage();
|
||||
manager.autoTagState.reviewIndex = 1;
|
||||
const second = manager.loadNextReviewImage();
|
||||
reply(1, { success: true, prompt_id: 2, tags: ['current'] });
|
||||
await second;
|
||||
reply(0, { success: true, prompt_id: 1, tags: ['stale'] });
|
||||
await first;
|
||||
assert.deepEqual(Array.from(manager.autoTagState.currentTags), ['current']);
|
||||
const third = manager.loadNextReviewImage();
|
||||
manager.cancelReview();
|
||||
reply(2, { success: true, prompt_id: 2, tags: ['closed'] });
|
||||
await third;
|
||||
assert.equal(manager.autoTagState.currentTags.length, 0);
|
||||
});
|
||||
@@ -0,0 +1,30 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const vm = require('node:vm');
|
||||
const path = require('node:path');
|
||||
const ImageHelpers = require('../../web/js/image-helpers.js');
|
||||
|
||||
test('gallery metadata requests retain root identity without including cache parameters in filenames', async () => {
|
||||
const requests = [];
|
||||
const context = vm.createContext({ document: { addEventListener() {} }, window: {}, console,
|
||||
ImageHelpers, ApiClient: { fetchUrl: async url => {
|
||||
requests.push(url);
|
||||
return { ok: true, json: async () => ({ success: true, prompt: { text: url } }) };
|
||||
} },
|
||||
});
|
||||
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
|
||||
const manager = Object.create(Manager.prototype);
|
||||
manager.extractImageMetadata = async () => null;
|
||||
manager.displayMetadata = (metadata, panel) => { panel.prompt = metadata.positivePrompt; };
|
||||
for (const root of [0, 1]) {
|
||||
const panel = {};
|
||||
const expected = `/prompt_manager/images/prompt/sub%2Fsame.png?root=${root}`;
|
||||
await manager.loadImageMetadata(`/prompt_manager/images/serve/sub/same.png?v=123&root=${root}`, panel);
|
||||
assert.equal(panel.prompt, expected);
|
||||
}
|
||||
assert.deepEqual(requests, [
|
||||
'/prompt_manager/images/prompt/sub%2Fsame.png?root=0',
|
||||
'/prompt_manager/images/prompt/sub%2Fsame.png?root=1',
|
||||
]);
|
||||
});
|
||||
@@ -0,0 +1,63 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const {
|
||||
DEFAULT_SETTINGS,
|
||||
normalizeGallerySettings,
|
||||
gridColumnsStyle,
|
||||
LIMIT_OPTIONS,
|
||||
} = require("../../web/js/gallery-settings.js");
|
||||
|
||||
test("defaults are a complete, sane settings object", () => {
|
||||
const s = normalizeGallerySettings(null);
|
||||
assert.deepEqual(s, DEFAULT_SETTINGS);
|
||||
assert.equal(s.defaultViewMode, "grid");
|
||||
assert.equal(s.defaultLimit, 100);
|
||||
assert.equal(s.gridColumns, 8);
|
||||
assert.equal(s.showImageInfo, true);
|
||||
assert.equal(s.sidebarCollapsedByDefault, false);
|
||||
});
|
||||
|
||||
test("stored values override defaults and unknown keys are kept", () => {
|
||||
const s = normalizeGallerySettings({ defaultViewMode: "list", defaultLimit: 500, gridColumns: 10, showImageInfo: false, videoLoop: false, extra: 1 });
|
||||
assert.equal(s.defaultViewMode, "list");
|
||||
assert.equal(s.defaultLimit, 500);
|
||||
assert.equal(s.gridColumns, 10);
|
||||
assert.equal(s.showImageInfo, false);
|
||||
assert.equal(s.videoLoop, false);
|
||||
assert.equal(s.extra, 1);
|
||||
});
|
||||
|
||||
test("invalid stored values fall back to defaults instead of breaking the page", () => {
|
||||
const s = normalizeGallerySettings({ defaultViewMode: "mosaic", defaultLimit: 7, gridColumns: "lots", showImageInfo: "no" });
|
||||
assert.equal(s.defaultViewMode, "grid");
|
||||
assert.equal(s.defaultLimit, 100);
|
||||
assert.equal(s.gridColumns, 8);
|
||||
assert.equal(s.showImageInfo, true);
|
||||
});
|
||||
|
||||
test("numeric strings from <select> values are accepted and columns are clamped", () => {
|
||||
assert.equal(normalizeGallerySettings({ defaultLimit: "200" }).defaultLimit, 200);
|
||||
assert.equal(normalizeGallerySettings({ gridColumns: "10" }).gridColumns, 10);
|
||||
assert.equal(normalizeGallerySettings({ gridColumns: 40 }).gridColumns, 12);
|
||||
assert.equal(normalizeGallerySettings({ gridColumns: 1 }).gridColumns, 2);
|
||||
assert.deepEqual(LIMIT_OPTIONS, [50, 100, 200, 500]);
|
||||
});
|
||||
|
||||
test("normalize never mutates its input", () => {
|
||||
const raw = { gridColumns: 4 };
|
||||
normalizeGallerySettings(raw);
|
||||
assert.deepEqual(raw, { gridColumns: 4 });
|
||||
});
|
||||
|
||||
test("gridColumnsStyle produces an explicit grid template", () => {
|
||||
assert.equal(gridColumnsStyle(10), "repeat(10, minmax(0, 1fr))");
|
||||
assert.equal(gridColumnsStyle("6"), "repeat(6, minmax(0, 1fr))");
|
||||
assert.equal(gridColumnsStyle("x"), "repeat(8, minmax(0, 1fr))");
|
||||
});
|
||||
|
||||
test("infinite scrolling is off by default and must be a boolean", () => {
|
||||
assert.equal(DEFAULT_SETTINGS.infiniteScroll, false);
|
||||
assert.equal(normalizeGallerySettings({ infiniteScroll: true }).infiniteScroll, true);
|
||||
assert.equal(normalizeGallerySettings({ infiniteScroll: "yes" }).infiniteScroll, false);
|
||||
});
|
||||
@@ -0,0 +1,90 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const path = require('node:path');
|
||||
const vm = require('node:vm');
|
||||
|
||||
function fixture(status, settings = {}) {
|
||||
const timers = [];
|
||||
const prompts = [];
|
||||
const requests = [];
|
||||
const buttons = {};
|
||||
const storage = new Map();
|
||||
const document = {
|
||||
addEventListener() {},
|
||||
createElement: () => ({ addEventListener() {} }),
|
||||
body: { appendChild() {}, removeChild() {} },
|
||||
getElementById: id => buttons[id] ||= { addEventListener(type, handler) { this[type] = handler; } },
|
||||
};
|
||||
const context = vm.createContext({ document, window: {}, console: { debug() {}, warn() {} },
|
||||
setTimeout: callback => timers.push(callback),
|
||||
localStorage: { setItem: (key, value) => storage.set(key, value) },
|
||||
ApiClient: { fetchUrl: async () => ({ json: async () => ({ success: true,
|
||||
total: 10000, images: Array.from({ length: 50 }, () => ({ thumbnail_url: null })),
|
||||
}) }) },
|
||||
});
|
||||
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
|
||||
const manager = Object.create(Manager.prototype);
|
||||
Object.assign(manager, { getSettings: () => ({ checkThumbnailsAtStartup: true, ...settings }),
|
||||
api: { get: async url => { requests.push(url); return status; } },
|
||||
});
|
||||
const originalPrompt = manager.showThumbnailGenerationPrompt.bind(manager);
|
||||
manager.showThumbnailGenerationPrompt = (...args) => prompts.push(args);
|
||||
async function check() {
|
||||
await manager.checkThumbnailsAtStartup();
|
||||
await timers.shift()();
|
||||
}
|
||||
return { manager, check, prompts, requests, settings, buttons, storage, originalPrompt };
|
||||
}
|
||||
|
||||
test('startup uses exact collection counts instead of extrapolating the newest 50 images', async () => {
|
||||
const f = fixture({ success: true, total_images: 10000, pending: 50, ready: 9950, collection_id: 'a' });
|
||||
await f.check();
|
||||
assert.deepEqual(f.requests, ['/prompt_manager/images/thumbnails/status']);
|
||||
assert.equal(f.prompts.length, 0);
|
||||
});
|
||||
|
||||
test('startup prompt reports the actual number needing generation', async () => {
|
||||
const f = fixture({ success: true, total_images: 10000, pending: 1200, collection_id: 'a' });
|
||||
await f.check();
|
||||
assert.equal(f.prompts.length, 1);
|
||||
assert.deepEqual(f.prompts[0].slice(0, 2), [10000, 1200]);
|
||||
});
|
||||
|
||||
test('a complete collection never prompts again', async () => {
|
||||
const f = fixture({ success: true, total_images: 10000, pending: 0, collection_id: 'a' });
|
||||
await f.check();
|
||||
await f.check();
|
||||
assert.equal(f.prompts.length, 0);
|
||||
});
|
||||
|
||||
test('Not Now persists the collection state and suppresses the unchanged prompt', async () => {
|
||||
const status = { success: true, total_images: 100, pending: 30, collection_id: 'a' };
|
||||
const f = fixture(status);
|
||||
await f.check();
|
||||
f.originalPrompt(...f.prompts[0]);
|
||||
f.buttons.thumbnailPromptNo.click();
|
||||
Object.assign(f.settings, JSON.parse(f.storage.get('gallerySettings')));
|
||||
f.prompts.length = 0;
|
||||
await f.check();
|
||||
assert.equal(f.prompts.length, 0);
|
||||
status.pending = 40;
|
||||
await f.check();
|
||||
assert.equal(f.prompts.length, 1);
|
||||
});
|
||||
|
||||
test('disabled checks do not request collection status', async () => {
|
||||
const f = fixture({}, { checkThumbnailsAtStartup: false });
|
||||
await f.check();
|
||||
assert.equal(f.requests.length, 0);
|
||||
assert.equal(f.prompts.length, 0);
|
||||
});
|
||||
|
||||
test('saving gallery settings retains the remembered thumbnail dismissal', () => {
|
||||
const f = fixture({}, { thumbnailPromptDismissed: 'remembered-state' });
|
||||
f.manager.showNotification = () => {};
|
||||
f.manager.hideSettings = () => {};
|
||||
f.manager.applySettings = () => {};
|
||||
f.manager.saveSettings();
|
||||
assert.equal(JSON.parse(f.storage.get('gallerySettings')).thumbnailPromptDismissed, 'remembered-state');
|
||||
});
|
||||
@@ -0,0 +1,56 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { escapeHtml, escapeAttr } = require("../../web/js/html-escape.js");
|
||||
|
||||
test("escapes the five HTML-significant characters", () => {
|
||||
assert.equal(escapeHtml(`&<>"'`), "&<>"'");
|
||||
});
|
||||
|
||||
test("a double quote cannot break out of an attribute", () => {
|
||||
const payload = `" onmouseover="alert(1)`;
|
||||
const escaped = escapeHtml(payload);
|
||||
assert.ok(!escaped.includes(`"`), escaped);
|
||||
assert.equal(escaped, `" onmouseover="alert(1)`);
|
||||
assert.equal(`<a data-tag="${escaped}">`, `<a data-tag="" onmouseover="alert(1)">`);
|
||||
});
|
||||
|
||||
test("a single quote cannot break out of a single-quoted attribute", () => {
|
||||
const escaped = escapeHtml(`' onclick='alert(1)`);
|
||||
assert.ok(!escaped.includes(`'`), escaped);
|
||||
assert.equal(escaped, `' onclick='alert(1)`);
|
||||
});
|
||||
|
||||
test("a script tag is neutralised in text context", () => {
|
||||
assert.equal(escapeHtml("<script>x</script>"), "<script>x</script>");
|
||||
});
|
||||
|
||||
test("unicode and ordinary text pass through unchanged", () => {
|
||||
assert.equal(escapeHtml("héllo wörld 日本語 🙂"), "héllo wörld 日本語 🙂");
|
||||
assert.equal(escapeHtml("plain text, no specials"), "plain text, no specials");
|
||||
});
|
||||
|
||||
test("already escaped text is escaped again, never double-decoded", () => {
|
||||
assert.equal(escapeHtml("&"), "&amp;");
|
||||
});
|
||||
|
||||
test("null and undefined become the empty string", () => {
|
||||
assert.equal(escapeHtml(null), "");
|
||||
assert.equal(escapeHtml(undefined), "");
|
||||
});
|
||||
|
||||
test("numbers and booleans are coerced to strings", () => {
|
||||
assert.equal(escapeHtml(42), "42");
|
||||
assert.equal(escapeHtml(0), "0");
|
||||
assert.equal(escapeHtml(1.5), "1.5");
|
||||
assert.equal(escapeHtml(false), "false");
|
||||
});
|
||||
|
||||
test("objects are coerced through String()", () => {
|
||||
assert.equal(escapeHtml({ toString: () => "<b>" }), "<b>");
|
||||
});
|
||||
|
||||
test("escapeAttr is the same function as escapeHtml", () => {
|
||||
assert.equal(escapeAttr, escapeHtml);
|
||||
assert.equal(escapeAttr(`a"b`), "a"b");
|
||||
});
|
||||
@@ -0,0 +1,123 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const {
|
||||
formatImageCaption,
|
||||
workflowDocumentHtml,
|
||||
autotagSingleBody,
|
||||
fallbackImageSource,
|
||||
} = require("../../web/js/image-helpers.js");
|
||||
|
||||
const WD14 = { modelType: "wd14-vit", generalThreshold: 0.35, characterThreshold: 0.85, prompt: "ignored" };
|
||||
const CAPTION = { modelType: "blip", generalThreshold: 0.35, characterThreshold: 0.85, prompt: "describe it" };
|
||||
|
||||
test("autotagSingleBody sends image_id for a database row", () => {
|
||||
const body = autotagSingleBody({ id: 42, prompt_id: 7, image_path: "/out/a.png" }, WD14);
|
||||
assert.equal(body.image_id, 42);
|
||||
assert.equal("path" in body, false);
|
||||
});
|
||||
|
||||
test("autotagSingleBody sends path for an output-folder entry whose id is a digest", () => {
|
||||
const body = autotagSingleBody({ id: "3f2a9c1b0d4e5f67", path: "/out/sub/a.png", filename: "a.png" }, WD14);
|
||||
assert.equal(body.path, "/out/sub/a.png");
|
||||
assert.equal("image_id" in body, false);
|
||||
});
|
||||
|
||||
test("autotagSingleBody prefers path when a numeric-looking string id has no prompt link", () => {
|
||||
const body = autotagSingleBody({ id: "1234567890abcdef", path: "/out/b.png" }, WD14);
|
||||
assert.deepEqual(Object.keys(body).filter((k) => k === "path" || k === "image_id"), ["path"]);
|
||||
});
|
||||
|
||||
test("autotagSingleBody adds wd14 thresholds and no prompt", () => {
|
||||
const body = autotagSingleBody({ id: 1 }, WD14);
|
||||
assert.equal(body.model_type, "wd14-vit");
|
||||
assert.equal(body.general_threshold, 0.35);
|
||||
assert.equal(body.character_threshold, 0.85);
|
||||
assert.equal("prompt" in body, false);
|
||||
});
|
||||
|
||||
test("autotagSingleBody adds the prompt for caption models and no thresholds", () => {
|
||||
const body = autotagSingleBody({ id: 1 }, CAPTION);
|
||||
assert.equal(body.model_type, "blip");
|
||||
assert.equal(body.prompt, "describe it");
|
||||
assert.equal("general_threshold" in body, false);
|
||||
});
|
||||
|
||||
test("autotagSingleBody does not mutate its inputs", () => {
|
||||
const image = Object.freeze({ id: 1 });
|
||||
const settings = Object.freeze({ ...WD14 });
|
||||
assert.doesNotThrow(() => autotagSingleBody(image, settings));
|
||||
});
|
||||
|
||||
const XSS = '"><img src=x onerror=1>';
|
||||
const formatFileSize = (bytes) => `${bytes} B`;
|
||||
|
||||
test("formatImageCaption joins dimensions and size with the separator", () => {
|
||||
const caption = formatImageCaption({ width: 640, height: 480, file_size: 2048 }, { formatFileSize });
|
||||
assert.equal(caption, "640×480 | 2048 B");
|
||||
});
|
||||
|
||||
test("formatImageCaption honours a custom separator and omits a missing size", () => {
|
||||
assert.equal(formatImageCaption({ width: 1, height: 2 }, { formatFileSize, separator: " • " }), "1×2");
|
||||
});
|
||||
|
||||
test("formatImageCaption falls back to Unknown size when a dimension is missing", () => {
|
||||
assert.equal(formatImageCaption({ width: 640 }, { formatFileSize }), "Unknown size");
|
||||
assert.equal(formatImageCaption({}, { formatFileSize }), "Unknown size");
|
||||
assert.equal(formatImageCaption(null, { formatFileSize }), "Unknown size");
|
||||
});
|
||||
|
||||
test("formatImageCaption escapes string dimensions so they cannot break out of an attribute", () => {
|
||||
const caption = formatImageCaption({ width: XSS, height: 10 }, { formatFileSize });
|
||||
assert.ok(!caption.includes("<"), caption);
|
||||
assert.ok(!caption.includes('"'), caption);
|
||||
assert.ok(caption.includes(""><img"), caption);
|
||||
});
|
||||
|
||||
test("formatImageCaption ignores a non-numeric file_size instead of formatting it", () => {
|
||||
const caption = formatImageCaption({ width: 1, height: 1, file_size: XSS }, { formatFileSize });
|
||||
assert.equal(caption, "1×1");
|
||||
});
|
||||
|
||||
test("formatImageCaption escapes whatever formatFileSize returns", () => {
|
||||
const caption = formatImageCaption({ width: 1, height: 1, file_size: 5 }, { formatFileSize: () => "<b>" });
|
||||
assert.equal(caption, "1×1 | <b>");
|
||||
});
|
||||
|
||||
test("workflowDocumentHtml escapes the JSON so it cannot close the pre block", () => {
|
||||
const html = workflowDocumentHtml({ note: "</pre><script>alert(1)</script>" });
|
||||
assert.ok(html.includes("</pre><script>"), html);
|
||||
assert.ok(!html.includes("</pre><script>"), html);
|
||||
assert.ok(html.startsWith("<!DOCTYPE html>"));
|
||||
assert.ok(html.includes("<title>ComfyUI Workflow Data</title>"));
|
||||
});
|
||||
|
||||
test("workflowDocumentHtml pretty-prints the workflow with two-space indentation", () => {
|
||||
const html = workflowDocumentHtml({ a: 1 });
|
||||
assert.ok(html.includes('{\n "a": 1\n}'), html);
|
||||
});
|
||||
|
||||
test("fallbackImageSource offers the original once when a thumbnail fails", () => {
|
||||
const dataset = { original: "/serve/a.png", thumbnail: "/serve/thumbnails/a_thumb.png?v=1" };
|
||||
assert.equal(fallbackImageSource("/serve/thumbnails/a_thumb.png?v=1", dataset), "/serve/a.png");
|
||||
});
|
||||
|
||||
test("fallbackImageSource gives up when the original itself failed or was already tried", () => {
|
||||
assert.equal(fallbackImageSource("/serve/a.png", { original: "/serve/a.png" }), null);
|
||||
assert.equal(fallbackImageSource("/serve/t.png", { original: "/serve/a.png", fellBack: "1" }), null);
|
||||
assert.equal(fallbackImageSource("/serve/t.png", {}), null);
|
||||
assert.equal(fallbackImageSource("/serve/t.png", undefined), null);
|
||||
});
|
||||
|
||||
const { metadataLookupUrl } = require('../../web/js/image-helpers.js');
|
||||
|
||||
test('metadata lookup keeps root identity and strips thumbnail cache query parameters', () => {
|
||||
assert.equal(metadataLookupUrl('/prompt_manager/images/serve/sub/a.png?root=1&v=123'), '/prompt_manager/images/prompt/sub%2Fa.png?root=1');
|
||||
assert.equal(metadataLookupUrl('/prompt_manager/images/serve/a.png?root=0'), '/prompt_manager/images/prompt/a.png?root=0');
|
||||
});
|
||||
|
||||
test('metadata lookup preserves encoded filename characters and works behind a proxy', () => {
|
||||
assert.equal(metadataLookupUrl('https://example.test/ComfyBackendDirect/prompt_manager/images/serve/space%20%23%3F%2520%20caf%C3%A9.png?root=1'), '/prompt_manager/images/prompt/space%20%23%3F%2520%20caf%C3%A9.png?root=1');
|
||||
assert.equal(metadataLookupUrl('/prompt_manager/images/serve/a.png?v=123'), '/prompt_manager/images/prompt/a.png');
|
||||
assert.equal(metadataLookupUrl('/elsewhere/a.png'), null);
|
||||
});
|
||||
@@ -0,0 +1,65 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { nextPageState, shouldSaveEdit, hasMorePages, shouldLoadMore } = require("../../web/js/list-state.js");
|
||||
|
||||
test("hasMorePages follows pagination.has_more only while pages keep arriving", () => {
|
||||
assert.equal(hasMorePages({ pagination: { has_more: true } }, 5000), true);
|
||||
assert.equal(hasMorePages({ pagination: { has_more: false } }, 5000), false);
|
||||
assert.equal(hasMorePages({ pagination: { has_more: true } }, 0), false, "an empty page ends the loop");
|
||||
assert.equal(hasMorePages({}, 10), false, "no pagination block means a single page");
|
||||
assert.equal(hasMorePages(null, 10), false);
|
||||
});
|
||||
|
||||
const current = { page: 2, limit: 50, total: 120, totalPages: 3 };
|
||||
|
||||
test("nextPageState moves to the requested page once the fetch succeeded", () => {
|
||||
assert.deepEqual(nextPageState(current, { page: 3, total: 120 }), { page: 3, limit: 50, total: 120, totalPages: 3 });
|
||||
});
|
||||
|
||||
test("nextPageState clamps a page beyond the end when the total shrank", () => {
|
||||
assert.deepEqual(nextPageState(current, { page: 3, total: 60 }), { page: 2, limit: 50, total: 60, totalPages: 2 });
|
||||
});
|
||||
|
||||
test("nextPageState never goes below page 1 and has at least one page", () => {
|
||||
assert.deepEqual(nextPageState(current, { page: 0, total: 0 }), { page: 1, limit: 50, total: 0, totalPages: 1 });
|
||||
assert.equal(nextPageState(current, { page: -5, total: 10 }).page, 1);
|
||||
});
|
||||
|
||||
test("nextPageState keeps the current page and total when the response omits them", () => {
|
||||
assert.deepEqual(nextPageState(current, {}), current);
|
||||
assert.deepEqual(nextPageState(current, { total: "many", page: "x" }), current);
|
||||
});
|
||||
|
||||
test("nextPageState accepts a new limit and recomputes totalPages", () => {
|
||||
assert.deepEqual(nextPageState(current, { page: 1, limit: 25, total: 120 }), { page: 1, limit: 25, total: 120, totalPages: 5 });
|
||||
assert.equal(nextPageState(current, { limit: 0 }).limit, 50, "a zero limit is ignored");
|
||||
});
|
||||
|
||||
test("nextPageState does not mutate its input", () => {
|
||||
const before = { ...current };
|
||||
nextPageState(current, { page: 1, total: 5 });
|
||||
assert.deepEqual(current, before);
|
||||
});
|
||||
|
||||
test("shouldSaveEdit is true only for a non-empty text that changed", () => {
|
||||
assert.equal(shouldSaveEdit("old", "new"), true);
|
||||
assert.equal(shouldSaveEdit("same", "same"), false);
|
||||
assert.equal(shouldSaveEdit("same", " same \n"), false, "whitespace-only changes are not saved");
|
||||
assert.equal(shouldSaveEdit("old", " "), false, "an emptied prompt is not saved");
|
||||
assert.equal(shouldSaveEdit("old", null), false);
|
||||
});
|
||||
|
||||
test("shouldLoadMore asks for the next page only near the bottom with pages left", () => {
|
||||
const base = {
|
||||
infiniteScroll: true, loading: false, page: 1, limit: 100, total: 1000,
|
||||
scrollTop: 0, clientHeight: 800, scrollHeight: 4000, threshold: 600,
|
||||
};
|
||||
assert.equal(shouldLoadMore(base), false, "far from the bottom");
|
||||
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700 }), true, "within the threshold");
|
||||
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700, loading: true }), false, "a load is in flight");
|
||||
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700, infiniteScroll: false }), false, "setting off");
|
||||
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700, page: 10 }), false, "last page reached");
|
||||
assert.equal(shouldLoadMore({ ...base, scrollHeight: 700 }), true, "content shorter than the viewport");
|
||||
assert.equal(shouldLoadMore({ ...base, total: 0, scrollHeight: 700 }), false, "nothing to load");
|
||||
});
|
||||
@@ -0,0 +1,102 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { parsePngTextChunks, PNG_SIGNATURE } = require("../../web/js/png-metadata.js");
|
||||
|
||||
const SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
|
||||
|
||||
function u32(n) {
|
||||
return [(n >>> 24) & 0xff, (n >>> 16) & 0xff, (n >>> 8) & 0xff, n & 0xff];
|
||||
}
|
||||
|
||||
/** A chunk with a dummy CRC (the parser does not verify it). */
|
||||
function chunk(type, dataBytes) {
|
||||
return [...u32(dataBytes.length), ...Buffer.from(type, "latin1"), ...dataBytes, 0, 0, 0, 0];
|
||||
}
|
||||
|
||||
const IHDR = chunk("IHDR", [...u32(1), ...u32(1), 8, 6, 0, 0, 0]);
|
||||
const IEND = chunk("IEND", []);
|
||||
|
||||
function tEXt(keyword, text) {
|
||||
return chunk("tEXt", [...Buffer.from(keyword, "latin1"), 0, ...Buffer.from(text, "latin1")]);
|
||||
}
|
||||
|
||||
function iTXt(keyword, text, { compressed = false } = {}) {
|
||||
return chunk("iTXt", [
|
||||
...Buffer.from(keyword, "latin1"), 0,
|
||||
compressed ? 1 : 0, 0, // compression flag, compression method
|
||||
0, // empty language tag
|
||||
0, // empty translated keyword
|
||||
...Buffer.from(text, "utf8"),
|
||||
]);
|
||||
}
|
||||
|
||||
function png(...chunks) {
|
||||
return Uint8Array.from([...SIGNATURE, ...IHDR, ...chunks.flat(), ...IEND]);
|
||||
}
|
||||
|
||||
test("exports the PNG signature", () => {
|
||||
assert.deepEqual([...PNG_SIGNATURE], SIGNATURE);
|
||||
});
|
||||
|
||||
test("reads a tEXt chunk written by ComfyUI", () => {
|
||||
const bytes = png(tEXt("prompt", '{"a":1}'));
|
||||
assert.deepEqual(parsePngTextChunks(bytes), { prompt: '{"a":1}' });
|
||||
});
|
||||
|
||||
test("reads several text chunks and an uncompressed iTXt chunk with UTF-8 text", () => {
|
||||
const bytes = png(tEXt("prompt", '{"a":1}'), iTXt("workflow", '{"b":"日本語 🙂"}'), tEXt("Software", "ComfyUI"));
|
||||
assert.deepEqual(parsePngTextChunks(bytes), {
|
||||
prompt: '{"a":1}',
|
||||
workflow: '{"b":"日本語 🙂"}',
|
||||
Software: "ComfyUI",
|
||||
});
|
||||
});
|
||||
|
||||
test("tEXt is Latin-1, so byte 0xE9 is an e-acute", () => {
|
||||
const bytes = png(chunk("tEXt", [...Buffer.from("Comment", "latin1"), 0, 0xe9]));
|
||||
assert.deepEqual(parsePngTextChunks(bytes), { Comment: "é" });
|
||||
});
|
||||
|
||||
test("skips compressed iTXt and zTXt payloads instead of returning garbage", () => {
|
||||
const zTXt = chunk("zTXt", [...Buffer.from("Comment", "latin1"), 0, 0, 0x78, 0x9c, 1, 2, 3]);
|
||||
const bytes = png(iTXt("workflow", "zzz", { compressed: true }), zTXt, tEXt("prompt", "ok"));
|
||||
assert.deepEqual(parsePngTextChunks(bytes), { prompt: "ok" });
|
||||
});
|
||||
|
||||
test("ignores non-text chunks and a text chunk without a NUL separator", () => {
|
||||
const bytes = png(chunk("IDAT", [1, 2, 3]), chunk("tEXt", [...Buffer.from("nonul", "latin1")]), tEXt("k", "v"));
|
||||
assert.deepEqual(parsePngTextChunks(bytes), { k: "v" });
|
||||
});
|
||||
|
||||
test("accepts an ArrayBuffer as well as a Uint8Array", () => {
|
||||
const bytes = png(tEXt("prompt", "x"));
|
||||
assert.deepEqual(parsePngTextChunks(bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength)), { prompt: "x" });
|
||||
});
|
||||
|
||||
test("returns {} for garbage, empty and non-PNG input", () => {
|
||||
assert.deepEqual(parsePngTextChunks(new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])), {});
|
||||
assert.deepEqual(parsePngTextChunks(new Uint8Array(0)), {});
|
||||
assert.deepEqual(parsePngTextChunks(Uint8Array.from(Buffer.from("GIF89a....................", "latin1"))), {});
|
||||
assert.deepEqual(parsePngTextChunks(null), {});
|
||||
assert.deepEqual(parsePngTextChunks("not bytes"), {});
|
||||
});
|
||||
|
||||
test("a file truncated inside a text chunk yields only the chunks that were complete", () => {
|
||||
const full = png(tEXt("prompt", "first"), tEXt("workflow", "second-second-second"));
|
||||
const cutInsideSecond = full.subarray(0, full.length - IEND.length - 6);
|
||||
assert.deepEqual(parsePngTextChunks(cutInsideSecond), { prompt: "first" });
|
||||
const cutInsideFirst = full.subarray(0, SIGNATURE.length + IHDR.length + 10);
|
||||
assert.deepEqual(parsePngTextChunks(cutInsideFirst), {});
|
||||
});
|
||||
|
||||
test("a chunk whose declared length overruns the file is ignored", () => {
|
||||
const bogus = [...u32(0x7fffffff), ...Buffer.from("tEXt", "latin1"), ...Buffer.from("k\0v", "latin1")];
|
||||
const bytes = Uint8Array.from([...SIGNATURE, ...IHDR, ...tEXt("ok", "1"), ...bogus]);
|
||||
assert.deepEqual(parsePngTextChunks(bytes), { ok: "1" });
|
||||
});
|
||||
|
||||
test("stops at IEND even if bytes follow it", () => {
|
||||
const bytes = Uint8Array.from([...png(tEXt("a", "1")), ...tEXt("b", "2")]);
|
||||
assert.deepEqual(parsePngTextChunks(bytes), { a: "1" });
|
||||
});
|
||||
@@ -0,0 +1,66 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const sort = require("../../web/js/prompt-list-sort.js");
|
||||
|
||||
test("default sort is recently used and is a listed option", () => {
|
||||
assert.equal(sort.DEFAULT_SORT, "last_used_desc");
|
||||
assert.ok(sort.SORT_OPTIONS.some((o) => o.value === sort.DEFAULT_SORT));
|
||||
});
|
||||
|
||||
test("options include Recently used and Most used, each with a label", () => {
|
||||
const values = sort.SORT_OPTIONS.map((o) => o.value);
|
||||
assert.ok(values.includes("last_used_desc"));
|
||||
assert.ok(values.includes("run_count_desc"));
|
||||
for (const option of sort.SORT_OPTIONS) assert.ok(option.label.trim(), option.value);
|
||||
assert.equal(new Set(values).size, values.length, "values are unique");
|
||||
});
|
||||
|
||||
test("options cannot be mutated by callers", () => {
|
||||
assert.throws(() => sort.SORT_OPTIONS.push({ value: "x", label: "x" }));
|
||||
// Non-strict code fails silently on frozen objects, so check the value survived
|
||||
const first = sort.SORT_OPTIONS[0].value;
|
||||
sort.SORT_OPTIONS[0].value = "hacked";
|
||||
assert.equal(sort.SORT_OPTIONS[0].value, first);
|
||||
assert.ok(sort.SORT_OPTIONS.every(Object.isFrozen));
|
||||
});
|
||||
|
||||
test("normalizeSort keeps known values and falls back for anything else", () => {
|
||||
assert.equal(sort.normalizeSort("run_count_desc"), "run_count_desc");
|
||||
for (const bad of [undefined, null, "", "created_at; DROP TABLE prompts", 42, "LAST_USED_DESC"]) {
|
||||
assert.equal(sort.normalizeSort(bad), sort.DEFAULT_SORT, String(bad));
|
||||
}
|
||||
});
|
||||
|
||||
test("buildRecentUrl computes offset from page and always sends a valid sort", () => {
|
||||
assert.equal(
|
||||
sort.buildRecentUrl({ page: 3, limit: 25, sort: "rating_desc" }),
|
||||
"/prompt_manager/recent?limit=25&offset=50&page=3&sort=rating_desc",
|
||||
);
|
||||
assert.equal(
|
||||
sort.buildRecentUrl({ sort: "bogus" }),
|
||||
"/prompt_manager/recent?limit=50&offset=0&page=1&sort=last_used_desc",
|
||||
);
|
||||
});
|
||||
|
||||
test("withSort returns a copy with the sort set and leaves the input untouched", () => {
|
||||
const params = new URLSearchParams({ text: "castle", sort: "text_asc" });
|
||||
const result = sort.withSort(params, "run_count_desc");
|
||||
assert.equal(result.get("sort"), "run_count_desc");
|
||||
assert.equal(result.get("text"), "castle");
|
||||
assert.equal(params.get("sort"), "text_asc");
|
||||
assert.equal(sort.withSort(new URLSearchParams(), "nope").get("sort"), sort.DEFAULT_SORT);
|
||||
});
|
||||
|
||||
test("formatRunCount labels zero, one and many runs", () => {
|
||||
assert.equal(sort.formatRunCount(0), "Never run");
|
||||
assert.equal(sort.formatRunCount(1), "1 run");
|
||||
assert.equal(sort.formatRunCount(642), "642 runs");
|
||||
assert.equal(sort.formatRunCount(1234), "1,234 runs");
|
||||
});
|
||||
|
||||
test("formatRunCount hides the badge for missing or invalid counts", () => {
|
||||
for (const bad of [undefined, null, -1, 1.5, "3", NaN]) {
|
||||
assert.equal(sort.formatRunCount(bad), "", String(bad));
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,78 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const {
|
||||
initialState,
|
||||
reduceScanEvent,
|
||||
fromStatus,
|
||||
dismissAction,
|
||||
progressLabel,
|
||||
} = require("../../web/js/scan-progress.js");
|
||||
|
||||
test("initialState is idle with zero counts", () => {
|
||||
assert.deepEqual(initialState(), {
|
||||
status: "idle", progress: 0, statusText: "", processed: 0, found: 0, added: 0, linked: 0, message: "",
|
||||
});
|
||||
});
|
||||
|
||||
test("a progress event moves the scan to running and copies the counts", () => {
|
||||
const before = initialState();
|
||||
const after = reduceScanEvent(before, {
|
||||
type: "progress", progress: 42, status: "Processing file 84/200...", processed: 84, found: 12,
|
||||
});
|
||||
assert.equal(after.status, "running");
|
||||
assert.equal(after.progress, 42);
|
||||
assert.equal(after.statusText, "Processing file 84/200...");
|
||||
assert.equal(after.processed, 84);
|
||||
assert.equal(after.found, 12);
|
||||
assert.equal(before.status, "idle", "reducer must not mutate its input");
|
||||
});
|
||||
|
||||
test("a complete event finishes the scan with the final totals", () => {
|
||||
const running = reduceScanEvent(initialState(), { type: "progress", progress: 90, status: "x", processed: 9, found: 1 });
|
||||
const done = reduceScanEvent(running, { type: "complete", processed: 10, found: 3, added: 2, linked: 1 });
|
||||
assert.equal(done.status, "done");
|
||||
assert.equal(done.progress, 100);
|
||||
assert.deepEqual([done.processed, done.found, done.added, done.linked], [10, 3, 2, 1]);
|
||||
});
|
||||
|
||||
test("an error event fails the scan and keeps the server message", () => {
|
||||
const failed = reduceScanEvent(initialState(), { type: "error", message: "No output directories found." });
|
||||
assert.equal(failed.status, "failed");
|
||||
assert.equal(failed.message, "No output directories found.");
|
||||
});
|
||||
|
||||
test("unknown events leave the state untouched", () => {
|
||||
const state = reduceScanEvent(initialState(), { type: "progress", progress: 5, status: "s", processed: 1, found: 0 });
|
||||
assert.equal(reduceScanEvent(state, { type: "mystery" }), state);
|
||||
assert.equal(reduceScanEvent(state, null), state);
|
||||
});
|
||||
|
||||
test("fromStatus rebuilds a running scan from the status endpoint", () => {
|
||||
const state = fromStatus({
|
||||
running: true,
|
||||
last_event: { type: "progress", progress: 30, status: "Processing file 3/10...", processed: 3, found: 1 },
|
||||
});
|
||||
assert.equal(state.status, "running");
|
||||
assert.equal(state.progress, 30);
|
||||
assert.equal(state.processed, 3);
|
||||
});
|
||||
|
||||
test("fromStatus is idle when nothing is running, even if a last event exists", () => {
|
||||
assert.equal(fromStatus({ running: false, last_event: { type: "complete", processed: 10 } }).status, "idle");
|
||||
assert.equal(fromStatus({ running: false, last_event: null }).status, "idle");
|
||||
assert.equal(fromStatus(null).status, "idle");
|
||||
});
|
||||
|
||||
test("dismissing the modal minimizes a running scan and closes otherwise", () => {
|
||||
const running = reduceScanEvent(initialState(), { type: "progress", progress: 1, status: "s", processed: 0, found: 0 });
|
||||
assert.equal(dismissAction(running), "minimize");
|
||||
assert.equal(dismissAction(initialState()), "close");
|
||||
assert.equal(dismissAction(reduceScanEvent(running, { type: "complete" })), "close");
|
||||
});
|
||||
|
||||
test("progressLabel summarises a running scan for the minimized pill", () => {
|
||||
const running = reduceScanEvent(initialState(), { type: "progress", progress: 42, status: "s", processed: 84, found: 12 });
|
||||
assert.equal(progressLabel(running), "Scanning 42% · 84 files · 12 prompts");
|
||||
assert.equal(progressLabel(initialState()), "");
|
||||
});
|
||||
@@ -0,0 +1,234 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { parseSseChunk, connect } = require("../../web/js/sse-stream.js");
|
||||
|
||||
test("parseSseChunk splits complete events and keeps the unfinished tail", () => {
|
||||
const { events, rest } = parseSseChunk('data: {"a":1}\n\ndata: {"b":2}\n\ndata: {"c"');
|
||||
assert.deepEqual(events.map((e) => e.data), ['{"a":1}', '{"b":2}']);
|
||||
assert.equal(rest, 'data: {"c"');
|
||||
});
|
||||
|
||||
test("parseSseChunk joins multi-line data, reads event names and ignores comments", () => {
|
||||
const { events, rest } = parseSseChunk(": keep-alive\nevent: progress\ndata: line1\ndata: line2\nid: 7\n\n");
|
||||
assert.deepEqual(events, [{ event: "progress", data: "line1\nline2" }]);
|
||||
assert.equal(rest, "");
|
||||
});
|
||||
|
||||
test("parseSseChunk accepts CRLF line endings and a missing space after the colon", () => {
|
||||
const { events } = parseSseChunk("data:x\r\n\r\ndata: y\r\n\r\n");
|
||||
assert.deepEqual(events.map((e) => e.data), ["x", "y"]);
|
||||
});
|
||||
|
||||
test("parseSseChunk returns nothing for an empty or whitespace buffer", () => {
|
||||
assert.deepEqual(parseSseChunk(""), { events: [], rest: "" });
|
||||
assert.deepEqual(parseSseChunk("\n\n"), { events: [], rest: "" });
|
||||
});
|
||||
|
||||
function streamOf(chunks) {
|
||||
const encoder = new TextEncoder();
|
||||
return new ReadableStream({
|
||||
start(controller) {
|
||||
for (const c of chunks) controller.enqueue(encoder.encode(c));
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function until(check, timeoutMs = 500) {
|
||||
return new Promise((resolve, reject) => {
|
||||
const started = Date.now();
|
||||
(function poll() {
|
||||
if (check()) return resolve();
|
||||
if (Date.now() - started > timeoutMs) return reject(new Error("timed out"));
|
||||
setTimeout(poll, 2);
|
||||
})();
|
||||
});
|
||||
}
|
||||
|
||||
test("connect POSTs and delivers each event to onmessage, then onclose", async () => {
|
||||
const seen = [];
|
||||
let closed = false;
|
||||
const calls = [];
|
||||
const fetchImpl = async (url, init) => {
|
||||
calls.push({ url, init });
|
||||
return new Response(streamOf(['data: {"n":1}\n\ndata: {"n', '":2}\n\n']), { status: 200 });
|
||||
};
|
||||
const source = connect("/prompt_manager/autotag/start?x=1", { method: "POST", fetchImpl });
|
||||
source.onmessage = (e) => seen.push(JSON.parse(e.data).n);
|
||||
source.onclose = () => { closed = true; };
|
||||
await until(() => closed);
|
||||
assert.deepEqual(seen, [1, 2]);
|
||||
assert.equal(calls[0].init.method, "POST");
|
||||
assert.equal(calls[0].url, "/prompt_manager/autotag/start?x=1");
|
||||
assert.equal(source.readyState, source.CLOSED);
|
||||
});
|
||||
|
||||
test("connect keeps streaming requests under the current proxy mount", async () => {
|
||||
const previousApiClient = globalThis.ApiClient;
|
||||
const previousLocation = globalThis.location;
|
||||
globalThis.ApiClient = require("../../web/js/api-client.js");
|
||||
globalThis.location = { pathname: "/ComfyBackendDirect/prompt_manager/gallery" };
|
||||
let requestedUrl;
|
||||
try {
|
||||
const source = connect("/prompt_manager/images/generate-thumbnails/progress?quality=80", {
|
||||
method: "POST",
|
||||
fetchImpl: async (url) => {
|
||||
requestedUrl = url;
|
||||
return new Response(streamOf([]), { status: 200 });
|
||||
},
|
||||
});
|
||||
await until(() => source.readyState === source.CLOSED);
|
||||
assert.equal(requestedUrl, "/ComfyBackendDirect/prompt_manager/images/generate-thumbnails/progress?quality=80");
|
||||
} finally {
|
||||
globalThis.ApiClient = previousApiClient;
|
||||
globalThis.location = previousLocation;
|
||||
}
|
||||
});
|
||||
|
||||
test("connect reports a non-2xx response through onerror and never onmessage", async () => {
|
||||
let error = null;
|
||||
const fetchImpl = async () => new Response("nope", { status: 405 });
|
||||
const source = connect("/x", { fetchImpl });
|
||||
source.onmessage = () => assert.fail("no message expected");
|
||||
source.onerror = (err) => { error = err; };
|
||||
await until(() => error !== null);
|
||||
assert.match(error.message, /405/);
|
||||
assert.equal(source.readyState, source.CLOSED);
|
||||
});
|
||||
|
||||
test("connect reports a network failure through onerror", async () => {
|
||||
let error = null;
|
||||
const source = connect("/x", { fetchImpl: async () => { throw new TypeError("Failed to fetch"); } });
|
||||
source.onerror = (err) => { error = err; };
|
||||
await until(() => error !== null);
|
||||
assert.equal(error.message, "Failed to fetch");
|
||||
});
|
||||
|
||||
test("close() aborts the request and suppresses later callbacks", async () => {
|
||||
let aborted = false;
|
||||
let errors = 0;
|
||||
const fetchImpl = (url, init) =>
|
||||
new Promise((resolve, reject) => {
|
||||
init.signal.addEventListener("abort", () => {
|
||||
aborted = true;
|
||||
reject(new DOMException("aborted", "AbortError"));
|
||||
});
|
||||
resolve(new Response(new ReadableStream({ start() {} }), { status: 200 }));
|
||||
});
|
||||
const source = connect("/x", { fetchImpl });
|
||||
source.onerror = () => errors++;
|
||||
await until(() => source.readyState === source.OPEN);
|
||||
source.close();
|
||||
assert.equal(aborted, true);
|
||||
assert.equal(source.readyState, source.CLOSED);
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
assert.equal(errors, 0);
|
||||
});
|
||||
|
||||
test("connect forwards a request body and an Accept header", async () => {
|
||||
let init = null;
|
||||
const fetchImpl = async (_url, i) => { init = i; return new Response(streamOf([]), { status: 200 }); };
|
||||
const source = connect("/x", { method: "POST", body: '{"a":1}', fetchImpl });
|
||||
await until(() => source.readyState === source.CLOSED);
|
||||
assert.equal(init.body, '{"a":1}');
|
||||
assert.equal(init.headers.Accept, "text/event-stream");
|
||||
});
|
||||
|
||||
test("events arriving with no onmessage handler are dropped without error", async () => {
|
||||
const fetchImpl = async () => new Response(streamOf(["data: 1\n\n", "data: 2\n\n"]), { status: 200 });
|
||||
const source = connect("/x", { fetchImpl });
|
||||
let errors = 0;
|
||||
source.onerror = () => errors++;
|
||||
await until(() => source.readyState === source.CLOSED);
|
||||
assert.equal(errors, 0);
|
||||
});
|
||||
|
||||
test("close() during delivery stops further messages and skips onclose", async () => {
|
||||
const seen = [];
|
||||
let closed = false;
|
||||
const fetchImpl = async () => new Response(streamOf(["data: 1\n\ndata: 2\n\ndata: 3\n\n"]), { status: 200 });
|
||||
const source = connect("/x", { fetchImpl });
|
||||
source.onmessage = (e) => { seen.push(e.data); if (e.data === "1") source.close(); };
|
||||
source.onclose = () => { closed = true; };
|
||||
await until(() => source.readyState === source.CLOSED);
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
assert.deepEqual(seen, ["1"]);
|
||||
assert.equal(closed, false);
|
||||
});
|
||||
|
||||
test("close() before the response arrives discards it silently", async () => {
|
||||
let resolveFetch;
|
||||
let errors = 0;
|
||||
const fetchImpl = () => new Promise((resolve) => { resolveFetch = resolve; });
|
||||
const source = connect("/x", { fetchImpl });
|
||||
source.onerror = () => errors++;
|
||||
source.close();
|
||||
resolveFetch(new Response(streamOf(["data: late\n\n"]), { status: 200 }));
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
assert.equal(errors, 0);
|
||||
assert.equal(source.readyState, source.CLOSED);
|
||||
});
|
||||
|
||||
test("connect defaults to GET and the global fetch", async () => {
|
||||
const original = globalThis.fetch;
|
||||
let init = null;
|
||||
globalThis.fetch = async (_url, i) => { init = i; return new Response(streamOf([]), { status: 200 }); };
|
||||
try {
|
||||
const source = connect("/x");
|
||||
await until(() => source.readyState === source.CLOSED);
|
||||
assert.equal(init.method, "GET");
|
||||
} finally {
|
||||
globalThis.fetch = original;
|
||||
}
|
||||
});
|
||||
|
||||
test("addEventListener receives named events and unnamed ones go to 'message'", async () => {
|
||||
const named = [];
|
||||
const plain = [];
|
||||
let closed = false;
|
||||
const fetchImpl = async () =>
|
||||
new Response(
|
||||
streamOf(['event: start\ndata: {"total":3}\n\n', 'data: {"n":1}\n\n', 'event: complete\ndata: {"ok":true}\n\n']),
|
||||
{ status: 200 },
|
||||
);
|
||||
const source = connect("/prompt_manager/images/generate-thumbnails/progress", { method: "POST", fetchImpl });
|
||||
source.addEventListener("start", (e) => named.push(["start", JSON.parse(e.data).total]));
|
||||
source.addEventListener("complete", (e) => named.push(["complete", JSON.parse(e.data).ok]));
|
||||
source.addEventListener("message", (e) => plain.push(JSON.parse(e.data).n));
|
||||
source.onclose = () => { closed = true; };
|
||||
await until(() => closed);
|
||||
assert.deepEqual(named, [["start", 3], ["complete", true]]);
|
||||
assert.deepEqual(plain, [1]);
|
||||
});
|
||||
|
||||
test("removeEventListener stops delivery and onmessage still sees every event", async () => {
|
||||
const all = [];
|
||||
const progress = [];
|
||||
let closed = false;
|
||||
const fetchImpl = async () =>
|
||||
new Response(streamOf(["event: progress\ndata: 1\n\n", "event: progress\ndata: 2\n\n"]), { status: 200 });
|
||||
const source = connect("/x", { fetchImpl });
|
||||
const listener = (e) => {
|
||||
progress.push(e.data);
|
||||
source.removeEventListener("progress", listener);
|
||||
};
|
||||
source.addEventListener("progress", listener);
|
||||
source.onmessage = (e) => all.push(`${e.event}:${e.data}`);
|
||||
source.onclose = () => { closed = true; };
|
||||
await until(() => closed);
|
||||
assert.deepEqual(progress, ["1"]);
|
||||
assert.deepEqual(all, ["progress:1", "progress:2"]);
|
||||
});
|
||||
|
||||
test("a listener that closes the source stops later events", async () => {
|
||||
const seen = [];
|
||||
const fetchImpl = async () =>
|
||||
new Response(streamOf(["event: error\ndata: boom\n\n", "event: progress\ndata: 2\n\n"]), { status: 200 });
|
||||
const source = connect("/x", { fetchImpl });
|
||||
source.addEventListener("error", (e) => { seen.push(e.data); source.close(); });
|
||||
source.addEventListener("progress", (e) => seen.push(e.data));
|
||||
await until(() => source.readyState === source.CLOSED);
|
||||
await new Promise((r) => setTimeout(r, 10));
|
||||
assert.deepEqual(seen, ["boom"]);
|
||||
});
|
||||
@@ -0,0 +1,39 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const {
|
||||
splitTagInput,
|
||||
applySuggestion,
|
||||
suggestUrl,
|
||||
moveSelection,
|
||||
} = require("../../web/js/tag-autocomplete.js");
|
||||
|
||||
test("splitTagInput separates committed tags from the tag being typed", () => {
|
||||
assert.deepEqual(splitTagInput("asian, portrait, sm"), { committed: ["asian", "portrait"], current: "sm" });
|
||||
assert.deepEqual(splitTagInput("asian,"), { committed: ["asian"], current: "" });
|
||||
assert.deepEqual(splitTagInput("asian, "), { committed: ["asian"], current: "" });
|
||||
assert.deepEqual(splitTagInput("asi"), { committed: [], current: "asi" });
|
||||
assert.deepEqual(splitTagInput(""), { committed: [], current: "" });
|
||||
assert.deepEqual(splitTagInput(" asian ,, Asian , dog, "), { committed: ["asian", "dog"], current: "" });
|
||||
});
|
||||
|
||||
test("applySuggestion replaces the partial tag and leaves the cursor ready for the next one", () => {
|
||||
assert.equal(applySuggestion("asian, sm", "smile"), "asian, smile, ");
|
||||
assert.equal(applySuggestion("", "asian"), "asian, ");
|
||||
assert.equal(applySuggestion("asian, ", "portrait"), "asian, portrait, ");
|
||||
assert.equal(applySuggestion("asian, asi", "asian"), "asian, ");
|
||||
});
|
||||
|
||||
test("suggestUrl encodes the prefix and the committed tags", () => {
|
||||
assert.equal(suggestUrl("asian, sm"), "/prompt_manager/tags/suggest?q=sm&with=asian&limit=15");
|
||||
assert.equal(suggestUrl("a&b, c d, x", 5), "/prompt_manager/tags/suggest?q=x&with=a%26b%2Cc+d&limit=5");
|
||||
assert.equal(suggestUrl(""), "/prompt_manager/tags/suggest?q=&with=&limit=15");
|
||||
});
|
||||
|
||||
test("moveSelection wraps around the list and starts from nothing selected", () => {
|
||||
assert.equal(moveSelection(-1, 1, 3), 0);
|
||||
assert.equal(moveSelection(-1, -1, 3), 2);
|
||||
assert.equal(moveSelection(2, 1, 3), 0);
|
||||
assert.equal(moveSelection(0, -1, 3), 2);
|
||||
assert.equal(moveSelection(1, 1, 0), -1);
|
||||
});
|
||||
@@ -0,0 +1,71 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const path = require('node:path');
|
||||
const vm = require('node:vm');
|
||||
function fixture() {
|
||||
const requests = [];
|
||||
const grid = { replaceChildren() {}, appendChild() {} };
|
||||
const context = vm.createContext({ window: {}, URLSearchParams, console: { error() {} },
|
||||
requestAnimationFrame() {}, document: { getElementById: id => id === 'tagPromptsGrid' ? grid : null,
|
||||
querySelector: () => null, createElement: () => ({ appendChild() {} }) },
|
||||
});
|
||||
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/tags-page.js'), 'utf8') + '\nTagsPageManager;', context, { filename: path.join(__dirname, '../../web/js/tags-page.js') });
|
||||
const manager = Object.create(Manager.prototype);
|
||||
Object.assign(manager, { selectedTags: ['cats'], filterMode: 'and', promptLimit: 20,
|
||||
promptOffset: 0, isLoadingPrompts: false, prompts: [], hasMorePrompts: false,
|
||||
fetchJson(url) { return new Promise((resolve, reject) => requests.push({ url, resolve, reject })); },
|
||||
renderTagPrompts() {}, appendTagPrompts() {}, updatePromptCountBadge() {}, updateLoadMoreButton() {},
|
||||
});
|
||||
const reply = (index, text, more = false) => requests[index].resolve({ success: true,
|
||||
prompts: [{ text }], pagination: { has_more: more, total: more ? 40 : 1 } });
|
||||
return { manager, requests, reply };
|
||||
}
|
||||
test('changing a filter starts a new request and late results cannot overwrite it', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const old = manager.loadTagPrompts();
|
||||
manager.selectedTags = ['dogs'];
|
||||
const current = manager.loadTagPrompts();
|
||||
assert.equal(requests.length, 2);
|
||||
assert.match(requests[1].url, /dogs/);
|
||||
reply(1, 'dog prompt'); await current;
|
||||
reply(0, 'cat prompt', true); await old;
|
||||
assert.equal(manager.prompts[0].text, 'dog prompt');
|
||||
assert.equal(manager.promptTotal, 1);
|
||||
assert.equal(manager.hasMorePrompts, false);
|
||||
});
|
||||
test('clearing selection invalidates pending results', async () => {
|
||||
const { manager, reply } = fixture();
|
||||
const old = manager.loadTagPrompts();
|
||||
manager.selectedTags = [];
|
||||
manager.clearPrompts();
|
||||
reply(0, 'cat prompt'); await old;
|
||||
assert.equal(manager.prompts.length, 0);
|
||||
assert.equal(manager.promptTotal, 0);
|
||||
assert.equal(manager.isLoadingPrompts, false);
|
||||
});
|
||||
test('a stale pagination response cannot append rows from the previous filter', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
Object.assign(manager, { prompts: [{ text: 'cat first page' }], hasMorePrompts: true, promptOffset: 20 });
|
||||
const old = manager.loadMoreTagPrompts();
|
||||
manager.selectedTags = ['dogs'];
|
||||
const current = manager.loadTagPrompts();
|
||||
assert.equal(requests.length, 2);
|
||||
reply(1, 'dog prompt'); await current;
|
||||
reply(0, 'cat second page', true); await old;
|
||||
assert.equal(manager.prompts.length, 1);
|
||||
assert.equal(manager.prompts[0].text, 'dog prompt');
|
||||
assert.equal(manager.promptOffset, 20);
|
||||
});
|
||||
test('a stale failure cannot finish the new request or clear its loading state', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const old = manager.loadTagPrompts();
|
||||
manager.selectedTags = ['dogs'];
|
||||
const current = manager.loadTagPrompts();
|
||||
assert.equal(requests.length, 2);
|
||||
requests[0].reject(new Error('old failure')); await old;
|
||||
assert.equal(manager.isLoadingPrompts, true);
|
||||
reply(1, 'dog prompt'); await current;
|
||||
assert.equal(manager.isLoadingPrompts, false);
|
||||
assert.equal(manager.prompts[0].text, 'dog prompt');
|
||||
});
|
||||
@@ -0,0 +1,42 @@
|
||||
import test from 'node:test';
|
||||
import assert from 'node:assert/strict';
|
||||
import { createScope, themeFromColor } from '../../web/js/shared/toolkit/core.js';
|
||||
|
||||
test('host background luminance determines mode without assuming palette names', () => {
|
||||
for (const color of ['#fff', '#F4F4F6', 'rgb(240, 240, 240)']) assert.equal(themeFromColor(color), 'light');
|
||||
for (const color of ['#222', '#202020', 'rgb(30, 32, 34)']) assert.equal(themeFromColor(color), 'dark');
|
||||
assert.equal(themeFromColor(''), null);
|
||||
assert.equal(themeFromColor('transparent'), null);
|
||||
});
|
||||
|
||||
test('scope cancels requests and callbacks, removes listeners and disposes once', () => {
|
||||
const scope = createScope();
|
||||
const target = new EventTarget();
|
||||
let events = 0;
|
||||
let cleanups = 0;
|
||||
scope.listen(target, 'change', () => events++);
|
||||
scope.cleanup(() => cleanups++);
|
||||
target.dispatchEvent(new Event('change'));
|
||||
assert.equal(events, 1);
|
||||
scope.dispose();
|
||||
scope.dispose();
|
||||
target.dispatchEvent(new Event('change'));
|
||||
assert.equal(events, 1);
|
||||
assert.equal(cleanups, 1);
|
||||
assert.equal(scope.signal.aborted, true);
|
||||
});
|
||||
|
||||
test('debounced work runs once with the latest value and stops after disposal', async () => {
|
||||
const scope = createScope();
|
||||
const values = [];
|
||||
const change = scope.debounce(value => values.push(value), 1);
|
||||
change('old');
|
||||
change('new');
|
||||
await new Promise(resolve => setTimeout(resolve, 20));
|
||||
assert.deepEqual(values, ['new']);
|
||||
change('cancelled');
|
||||
scope.dispose();
|
||||
change('after removal');
|
||||
await new Promise(resolve => setTimeout(resolve, 20));
|
||||
assert.deepEqual(values, ['new']);
|
||||
});
|
||||
@@ -0,0 +1,30 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const { resolveView, VIEWS } = require("../../web/js/view-router.js");
|
||||
|
||||
test("the tags deep link and its sub-routes resolve to the tags view", () => {
|
||||
assert.equal(resolveView("#/tags"), "tags");
|
||||
assert.equal(resolveView("#/tags/"), "tags");
|
||||
assert.equal(resolveView("#/tags/landscape,portrait?mode=or"), "tags");
|
||||
});
|
||||
|
||||
test("an empty, bare or unknown hash resolves to the prompts dashboard", () => {
|
||||
assert.equal(resolveView(""), "prompts");
|
||||
assert.equal(resolveView("#"), "prompts");
|
||||
assert.equal(resolveView("#/"), "prompts");
|
||||
assert.equal(resolveView("#/other"), "prompts");
|
||||
assert.equal(resolveView("#/tagsX"), "prompts", "prefix must end at a path boundary");
|
||||
assert.equal(resolveView("#tags"), "prompts");
|
||||
});
|
||||
|
||||
test("non-string input resolves to the prompts dashboard", () => {
|
||||
assert.equal(resolveView(undefined), "prompts");
|
||||
assert.equal(resolveView(null), "prompts");
|
||||
assert.equal(resolveView(42), "prompts");
|
||||
});
|
||||
|
||||
test("every resolved value is a listed view", () => {
|
||||
assert.ok(Object.isFrozen(VIEWS));
|
||||
for (const hash of ["#/tags", "", "#/x"]) assert.ok(VIEWS.includes(resolveView(hash)));
|
||||
});
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Teardown guard: fail when a file about to be deleted is still open.
|
||||
|
||||
Windows refuses to delete an open file, so a sqlite connection that was not
|
||||
closed explicitly turns into a PermissionError in CI there. Linux would let
|
||||
the deletion through silently; this check makes the Linux run catch the same
|
||||
mistake. It inspects /proc/self/fd and is a no-op where that does not exist.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
|
||||
def open_handles_under(path):
|
||||
"""Basenames of this process's open files at *path* or below it."""
|
||||
fd_dir = "/proc/self/fd"
|
||||
if not os.path.isdir(fd_dir):
|
||||
return []
|
||||
real = os.path.realpath(str(path))
|
||||
hits = []
|
||||
for fd in os.listdir(fd_dir):
|
||||
try:
|
||||
target = os.readlink(os.path.join(fd_dir, fd))
|
||||
except OSError:
|
||||
continue
|
||||
if not target.startswith(os.sep):
|
||||
continue # sockets, pipes, anon inodes
|
||||
target = target.replace(" (deleted)", "")
|
||||
if (
|
||||
target == real
|
||||
or target.startswith(real + os.sep)
|
||||
or target.startswith(real + "-")
|
||||
):
|
||||
hits.append(os.path.basename(target))
|
||||
return sorted(hits)
|
||||
|
||||
|
||||
def assert_closed(testcase, *paths):
|
||||
"""Fail *testcase* if any of *paths* (files or directories) is still open."""
|
||||
hits = [hit for path in paths for hit in open_handles_under(path)]
|
||||
if hits:
|
||||
testcase.fail(
|
||||
"Files still open at teardown (Windows could not delete them): "
|
||||
+ ", ".join(hits)
|
||||
)
|
||||
@@ -0,0 +1,420 @@
|
||||
"""
|
||||
API endpoint integration tests for PromptManager.
|
||||
|
||||
Uses aiohttp test client to test request/response contracts
|
||||
without requiring a full ComfyUI server.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import AioHTTPTestCase
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from py.api import PromptManagerAPI
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
class APITestCase(AioHTTPTestCase):
|
||||
"""Base class that stands up a test aiohttp app with PromptManager routes."""
|
||||
|
||||
async def get_application(self):
|
||||
self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
|
||||
self._temp_db.close()
|
||||
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self._temp_db.name)
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# Close every thread's connection first: Windows refuses to delete
|
||||
# an open database file, and executor threads hold connections too.
|
||||
self.api.db.close_all()
|
||||
for path in (
|
||||
self._temp_db.name,
|
||||
self._temp_db.name + "-wal",
|
||||
self._temp_db.name + "-shm",
|
||||
):
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
# ── helpers ────────────────────────────────────────────────────────
|
||||
|
||||
def _save_prompt(self, text="Test prompt", category=None, tags=None, rating=None):
|
||||
return self.api.db.save_prompt(
|
||||
text=text,
|
||||
category=category,
|
||||
tags=tags or [],
|
||||
rating=rating,
|
||||
prompt_hash=generate_prompt_hash(text),
|
||||
)
|
||||
|
||||
|
||||
class TestHealthEndpoint(APITestCase):
|
||||
|
||||
async def test_test_route(self):
|
||||
resp = await self.client.request("GET", "/prompt_manager/test")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertIn("timestamp", data)
|
||||
|
||||
|
||||
class TestRecentPrompts(APITestCase):
|
||||
|
||||
async def test_recent_empty(self):
|
||||
resp = await self.client.request("GET", "/prompt_manager/recent")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["pagination"]["total"], 0)
|
||||
self.assertEqual(len(data["results"]), 0)
|
||||
|
||||
async def test_recent_with_data(self):
|
||||
self._save_prompt("First prompt")
|
||||
self._save_prompt("Second prompt")
|
||||
resp = await self.client.request("GET", "/prompt_manager/recent")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertEqual(data["pagination"]["total"], 2)
|
||||
self.assertEqual(len(data["results"]), 2)
|
||||
|
||||
async def test_recent_pagination(self):
|
||||
for i in range(15):
|
||||
self._save_prompt(f"Prompt {i:02d}")
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/recent?limit=5&offset=0"
|
||||
)
|
||||
data = await resp.json()
|
||||
self.assertEqual(len(data["results"]), 5)
|
||||
self.assertEqual(data["pagination"]["total"], 15)
|
||||
self.assertTrue(data["pagination"]["has_more"])
|
||||
|
||||
async def test_recent_last_page(self):
|
||||
for i in range(7):
|
||||
self._save_prompt(f"Prompt {i}")
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/recent?limit=5&offset=5"
|
||||
)
|
||||
data = await resp.json()
|
||||
self.assertEqual(len(data["results"]), 2)
|
||||
self.assertFalse(data["pagination"]["has_more"])
|
||||
|
||||
|
||||
class TestSearch(APITestCase):
|
||||
|
||||
async def test_search_by_text(self):
|
||||
self._save_prompt("Beautiful mountain landscape")
|
||||
self._save_prompt("City skyline at night")
|
||||
resp = await self.client.request("GET", "/prompt_manager/search?text=mountain")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(len(data["results"]), 1)
|
||||
self.assertIn("mountain", data["results"][0]["text"])
|
||||
|
||||
async def test_search_by_category(self):
|
||||
self._save_prompt("Nature scene", category="nature")
|
||||
self._save_prompt("Urban view", category="urban")
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/search?category=nature"
|
||||
)
|
||||
data = await resp.json()
|
||||
self.assertEqual(len(data["results"]), 1)
|
||||
|
||||
async def test_search_by_tag(self):
|
||||
self._save_prompt("Tagged prompt", tags=["landscape", "sunset"])
|
||||
self._save_prompt("Other prompt", tags=["portrait"])
|
||||
resp = await self.client.request("GET", "/prompt_manager/search?tags=landscape")
|
||||
data = await resp.json()
|
||||
self.assertEqual(len(data["results"]), 1)
|
||||
|
||||
async def test_search_empty_result(self):
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/search?text=nonexistent_xyz"
|
||||
)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(len(data["results"]), 0)
|
||||
|
||||
|
||||
class TestSaveAndDelete(APITestCase):
|
||||
|
||||
async def test_save_prompt(self):
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/save",
|
||||
json={
|
||||
"text": "New prompt via API",
|
||||
"category": "test",
|
||||
"tags": ["api", "test"],
|
||||
},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertIn("prompt_id", data)
|
||||
|
||||
async def test_save_empty_text(self):
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/save", json={"text": ""}
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
|
||||
async def test_save_with_null_optional_fields(self):
|
||||
"""Regression for #143: UI sends null for empty optional fields."""
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/save",
|
||||
json={"text": "Null notes prompt", "category": None, "notes": None},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
async def test_save_with_null_text(self):
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/save", json={"text": None}
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_delete_prompt(self):
|
||||
pid = self._save_prompt("To delete")
|
||||
resp = await self.client.request("DELETE", f"/prompt_manager/delete/{pid}")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
async def test_delete_nonexistent(self):
|
||||
resp = await self.client.request("DELETE", "/prompt_manager/delete/99999")
|
||||
self.assertEqual(resp.status, 404)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
|
||||
|
||||
class TestUpdatePrompt(APITestCase):
|
||||
|
||||
async def test_update_text(self):
|
||||
pid = self._save_prompt("Original text")
|
||||
resp = await self.client.request(
|
||||
"PUT",
|
||||
f"/prompt_manager/prompts/{pid}",
|
||||
json={"text": "Updated text"},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
async def test_update_rating(self):
|
||||
pid = self._save_prompt("Ratable")
|
||||
resp = await self.client.request(
|
||||
"PUT",
|
||||
f"/prompt_manager/prompts/{pid}/rating",
|
||||
json={"rating": 4},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
async def test_add_tag(self):
|
||||
pid = self._save_prompt("Taggable", tags=["old"])
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
f"/prompt_manager/prompts/{pid}/tags",
|
||||
json={"tag": "new_tag"},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
|
||||
class TestUsageSorting(APITestCase):
|
||||
"""Server-side sort for /recent and /search (3.2.4, #90.1)."""
|
||||
|
||||
def _use(self, prompt_id, times=1):
|
||||
for _ in range(times):
|
||||
self.api.db.record_prompt_use(prompt_id)
|
||||
|
||||
async def _ids(self, path):
|
||||
resp = await self.client.request("GET", path)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
return [p["id"] for p in data["results"]]
|
||||
|
||||
async def test_recent_sorted_by_last_used_bubbles_rerun_prompt(self):
|
||||
old = self._save_prompt("Old prompt")
|
||||
self._save_prompt("New prompt")
|
||||
self._use(old)
|
||||
ids = await self._ids("/prompt_manager/recent?limit=1&sort=last_used_desc")
|
||||
self.assertEqual(ids, [old])
|
||||
|
||||
async def test_recent_sorted_by_run_count(self):
|
||||
# Most-run prompt is the oldest, so newest-first would give the opposite order
|
||||
most = self._save_prompt("Run thrice")
|
||||
least = self._save_prompt("Run once")
|
||||
self._use(most, times=3)
|
||||
self._use(least)
|
||||
ids = await self._ids("/prompt_manager/recent?sort=run_count_desc")
|
||||
self.assertEqual(ids, [most, least])
|
||||
|
||||
async def test_recent_default_is_newest_first(self):
|
||||
first = self._save_prompt("First")
|
||||
second = self._save_prompt("Second")
|
||||
self._use(first)
|
||||
self.assertEqual(await self._ids("/prompt_manager/recent"), [second, first])
|
||||
|
||||
async def test_unknown_sort_falls_back_instead_of_erroring(self):
|
||||
self._save_prompt("Only")
|
||||
ids = await self._ids(
|
||||
"/prompt_manager/recent?sort=created_at;DROP%20TABLE%20prompts"
|
||||
)
|
||||
self.assertEqual(len(ids), 1)
|
||||
|
||||
async def test_search_respects_sort(self):
|
||||
a = self._save_prompt("castle at dawn")
|
||||
self._save_prompt("castle at dusk")
|
||||
self._use(a)
|
||||
ids = await self._ids("/prompt_manager/search?text=castle&sort=last_used_desc")
|
||||
self.assertEqual(ids[0], a)
|
||||
|
||||
async def test_results_include_usage_fields(self):
|
||||
pid = self._save_prompt("Tracked")
|
||||
self._use(pid, times=2)
|
||||
resp = await self.client.request("GET", "/prompt_manager/recent")
|
||||
prompt = (await resp.json())["results"][0]
|
||||
self.assertEqual(prompt["run_count"], 2)
|
||||
self.assertTrue(prompt["last_used_at"])
|
||||
|
||||
|
||||
class TestTags(APITestCase):
|
||||
|
||||
async def test_get_all_tags(self):
|
||||
self._save_prompt("P1", tags=["alpha", "beta"])
|
||||
self._save_prompt("P2", tags=["beta", "gamma"])
|
||||
resp = await self.client.request("GET", "/prompt_manager/tags")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertIn("tags", data)
|
||||
|
||||
async def test_get_tag_stats(self):
|
||||
self._save_prompt("P1", tags=["common", "rare"])
|
||||
self._save_prompt("P2", tags=["common"])
|
||||
resp = await self.client.request("GET", "/prompt_manager/tags/stats")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
async def test_rename_tag(self):
|
||||
self._save_prompt("P1", tags=["old_name"])
|
||||
resp = await self.client.request(
|
||||
"PUT",
|
||||
"/prompt_manager/tags/old_name",
|
||||
json={"new_name": "new_name"},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
async def test_delete_tag(self):
|
||||
self._save_prompt("P1", tags=["removable"])
|
||||
resp = await self.client.request("DELETE", "/prompt_manager/tags/removable")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
async def test_merge_tags(self):
|
||||
self._save_prompt("P1", tags=["target"])
|
||||
self._save_prompt("P2", tags=["source"])
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/tags/merge",
|
||||
json={"source_tags": ["source"], "target_tag": "target"},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
|
||||
class TestCategories(APITestCase):
|
||||
|
||||
async def test_get_categories(self):
|
||||
self._save_prompt("P1", category="nature")
|
||||
self._save_prompt("P2", category="urban")
|
||||
resp = await self.client.request("GET", "/prompt_manager/categories")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertIn("categories", data)
|
||||
|
||||
|
||||
class TestStats(APITestCase):
|
||||
|
||||
async def test_get_stats(self):
|
||||
self._save_prompt("P1", category="test", tags=["t1"], rating=4)
|
||||
resp = await self.client.request("GET", "/prompt_manager/stats")
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertIn("stats", data)
|
||||
|
||||
|
||||
class TestExport(APITestCase):
|
||||
|
||||
async def test_export_json(self):
|
||||
self._save_prompt("Exportable", category="test", tags=["export"])
|
||||
resp = await self.client.request("GET", "/prompt_manager/export?format=json")
|
||||
self.assertEqual(resp.status, 200)
|
||||
# Export returns raw JSON file download, not a success envelope
|
||||
self.assertIn("Content-Disposition", resp.headers)
|
||||
text = await resp.text()
|
||||
data = json.loads(text)
|
||||
self.assertIn("prompts", data)
|
||||
self.assertEqual(len(data["prompts"]), 1)
|
||||
|
||||
|
||||
class TestResponseEnvelope(APITestCase):
|
||||
"""Verify all responses follow the {success: bool, ...} envelope."""
|
||||
|
||||
async def test_success_responses_have_success_true(self):
|
||||
self._save_prompt("Envelope test")
|
||||
endpoints = [
|
||||
("GET", "/prompt_manager/test"),
|
||||
("GET", "/prompt_manager/recent"),
|
||||
("GET", "/prompt_manager/search"),
|
||||
("GET", "/prompt_manager/categories"),
|
||||
("GET", "/prompt_manager/tags"),
|
||||
("GET", "/prompt_manager/stats"),
|
||||
]
|
||||
for method, path in endpoints:
|
||||
resp = await self.client.request(method, path)
|
||||
data = await resp.json()
|
||||
self.assertIn("success", data, f"Missing 'success' key in {method} {path}")
|
||||
self.assertTrue(
|
||||
data["success"], f"Expected success=True for {method} {path}"
|
||||
)
|
||||
|
||||
async def test_error_responses_have_success_false(self):
|
||||
resp = await self.client.request("DELETE", "/prompt_manager/delete/99999")
|
||||
data = await resp.json()
|
||||
self.assertIn("success", data)
|
||||
self.assertFalse(data["success"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,252 @@
|
||||
"""Route tests for GET /prompt_manager/backup and POST /prompt_manager/restore.
|
||||
|
||||
The backup must include rows still sitting in the WAL, the restore must
|
||||
stream the upload to disk, reject anything that is not a healthy prompts
|
||||
database and leave the live database untouched when it does.
|
||||
"""
|
||||
|
||||
import glob
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
from database.models import PromptModel # noqa: E402
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed # noqa: E402
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed # noqa: E402
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import AioHTTPTestCase
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from py.api import PromptManagerAPI, _public_path
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
def _texts(path):
|
||||
conn = sqlite3.connect(path)
|
||||
try:
|
||||
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class BackupRestoreTestCase(AioHTTPTestCase):
|
||||
async def get_application(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.db_path = os.path.join(self.tmpdir, "prompts.db")
|
||||
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self.db_path)
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
PromptModel.close_all_instances() # restore may have swapped api.db
|
||||
assert_closed(self, self.tmpdir)
|
||||
for name in os.listdir(self.tmpdir):
|
||||
os.unlink(os.path.join(self.tmpdir, name))
|
||||
os.rmdir(self.tmpdir)
|
||||
|
||||
def _save(self, text):
|
||||
return self.api.db.save_prompt(
|
||||
text=text, prompt_hash=generate_prompt_hash(text)
|
||||
)
|
||||
|
||||
def _write(self, name, data):
|
||||
path = os.path.join(self.tmpdir, name)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(data)
|
||||
return path
|
||||
|
||||
def _make_db_file(self, texts, name="upload.db"):
|
||||
path = os.path.join(self.tmpdir, name)
|
||||
db = PromptDatabase(path)
|
||||
for text in texts:
|
||||
db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
db.close_all()
|
||||
return path
|
||||
|
||||
async def _restore(self, path, field="database_file"):
|
||||
with open(path, "rb") as fh:
|
||||
form = aiohttp.FormData()
|
||||
form.add_field(
|
||||
field,
|
||||
fh.read(),
|
||||
filename=os.path.basename(path),
|
||||
content_type="application/octet-stream",
|
||||
)
|
||||
return await self.client.post("/prompt_manager/restore", data=form)
|
||||
|
||||
|
||||
class TestBackupRoute(BackupRestoreTestCase):
|
||||
async def test_backup_download_contains_rows_still_in_the_wal(self):
|
||||
for n in range(3):
|
||||
self._save(f"backup me {n}")
|
||||
self.assertGreater(os.path.getsize(self.db_path + "-wal"), 0)
|
||||
created = []
|
||||
real_mkstemp = tempfile.mkstemp
|
||||
|
||||
def spy_mkstemp(*args, **kwargs):
|
||||
result = real_mkstemp(*args, **kwargs)
|
||||
created.append(result[1])
|
||||
return result
|
||||
|
||||
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
|
||||
resp = await self.client.get("/prompt_manager/backup")
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(resp.content_type, "application/octet-stream")
|
||||
disposition = resp.headers["Content-Disposition"]
|
||||
self.assertIn("attachment", disposition)
|
||||
self.assertIn("prompts_backup_", disposition)
|
||||
self.assertTrue(disposition.rstrip('"').endswith(".db"))
|
||||
|
||||
body = await resp.read()
|
||||
downloaded = self._write("downloaded.db", body)
|
||||
self.assertEqual(
|
||||
_texts(downloaded), ["backup me 0", "backup me 1", "backup me 2"]
|
||||
)
|
||||
# The server-side temp file is gone once the response is complete
|
||||
self.assertEqual(len(created), 1)
|
||||
self.assertFalse(os.path.exists(created[0]))
|
||||
|
||||
async def test_backup_of_missing_database_is_404(self):
|
||||
self.api.db.model.db_path = os.path.join(self.tmpdir, "nope.db")
|
||||
resp = await self.client.get("/prompt_manager/backup")
|
||||
self.assertEqual(resp.status, 404)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
|
||||
async def test_backup_failure_is_500(self):
|
||||
self._save("x")
|
||||
with patch.object(self.api.db.model, "backup_database", return_value=False):
|
||||
resp = await self.client.get("/prompt_manager/backup")
|
||||
self.assertEqual(resp.status, 500)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("backup", data["error"].lower())
|
||||
|
||||
|
||||
class TestRestoreRoute(BackupRestoreTestCase):
|
||||
async def test_restore_replaces_database_and_reports_backup(self):
|
||||
self._save("old prompt")
|
||||
old_db = self.api.db
|
||||
upload = self._make_db_file(["restored a", "restored b"])
|
||||
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["prompt_count"], 2)
|
||||
# The response names the safety backup without the server's layout
|
||||
backups = glob.glob(self.db_path + ".backup_*")
|
||||
self.assertEqual(len(backups), 1)
|
||||
self.assertEqual(data["backup_created"], _public_path(backups[0]))
|
||||
self.assertNotIn(self.tmpdir, data["backup_created"])
|
||||
self.assertEqual(_texts(backups[0]), ["old prompt"])
|
||||
self.assertIsNot(self.api.db, old_db)
|
||||
self.assertEqual(self.api.db.model.db_path, self.db_path)
|
||||
|
||||
texts = sorted(p["text"] for p in self.api.db.search_prompts())
|
||||
self.assertEqual(texts, ["restored a", "restored b"])
|
||||
|
||||
async def test_restore_leaves_no_temp_files_behind(self):
|
||||
upload = self._make_db_file(["restored"])
|
||||
created = []
|
||||
real_mkstemp = tempfile.mkstemp
|
||||
|
||||
def spy_mkstemp(*args, **kwargs):
|
||||
result = real_mkstemp(*args, **kwargs)
|
||||
created.append(result[1])
|
||||
return result
|
||||
|
||||
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(len(created), 1)
|
||||
for suffix in ("", "-wal", "-shm", "-journal"):
|
||||
self.assertFalse(os.path.exists(created[0] + suffix))
|
||||
|
||||
async def test_text_file_is_rejected_and_database_untouched(self):
|
||||
self._save("keep me")
|
||||
upload = self._write("notes.txt", b"this is not a database")
|
||||
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("SQLite", data["error"])
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
self.assertFalse(
|
||||
any(n.startswith("prompts.db.backup_") for n in os.listdir(self.tmpdir))
|
||||
)
|
||||
|
||||
async def test_database_without_prompts_table_is_rejected(self):
|
||||
self._save("keep me")
|
||||
path = os.path.join(self.tmpdir, "other.db")
|
||||
conn = sqlite3.connect(path)
|
||||
conn.execute("CREATE TABLE things (id INTEGER)")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
resp = await self._restore(path)
|
||||
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("prompts", data["error"])
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
|
||||
async def test_wrong_field_name_is_rejected(self):
|
||||
upload = self._make_db_file(["restored"])
|
||||
resp = await self._restore(upload, field="file")
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("database_file", data["error"])
|
||||
|
||||
async def test_empty_upload_is_rejected(self):
|
||||
upload = self._write("empty.db", b"")
|
||||
resp = await self._restore(upload)
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("empty", data["error"].lower())
|
||||
|
||||
async def test_upload_over_the_size_limit_is_rejected_while_streaming(self):
|
||||
self._save("keep me")
|
||||
self.api.restore_max_bytes = 1024
|
||||
upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
|
||||
|
||||
resp = await self._restore(upload)
|
||||
|
||||
self.assertEqual(resp.status, 400)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("too large", data["error"].lower())
|
||||
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
||||
|
||||
async def test_non_multipart_request_is_rejected(self):
|
||||
resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes")
|
||||
self.assertIn(resp.status, (400, 500))
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,809 @@
|
||||
"""Tests for the core API module (py/api/__init__.py): helpers and base routes."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
_mock_server = MagicMock()
|
||||
_mock_server.PromptServer.instance.routes = MagicMock()
|
||||
sys.modules.setdefault("server", _mock_server)
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed # noqa: E402
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed # noqa: E402
|
||||
from database.models import PromptModel # noqa: E402
|
||||
|
||||
from aiohttp import web # noqa: E402
|
||||
from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
|
||||
|
||||
import py.api as api_module # noqa: E402
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from py.api import PromptManagerAPI, _public_error, _public_path # noqa: E402
|
||||
from py.config import GalleryConfig # noqa: E402
|
||||
|
||||
EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
|
||||
|
||||
|
||||
class FolderPathsFixture(unittest.TestCase):
|
||||
"""Temp ComfyUI tree exposed through a stub folder_paths module."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, self.tmpdir, True)
|
||||
self.addCleanup(assert_closed, self, self.tmpdir) # just before the rmtree
|
||||
self.addCleanup(PromptModel.close_all_instances)
|
||||
self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
|
||||
self.output_dir = self.comfy_dir / "output"
|
||||
self.output_dir.mkdir(parents=True)
|
||||
|
||||
self._orig_folder_paths = sys.modules.get("folder_paths")
|
||||
self.addCleanup(self._restore_folder_paths)
|
||||
self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None)
|
||||
self.addCleanup(self._restore_extra_env)
|
||||
|
||||
self.install_folder_paths(base_path=str(self.comfy_dir))
|
||||
|
||||
def install_folder_paths(self, **attrs):
|
||||
attrs.setdefault("get_output_directory", lambda: str(self.output_dir))
|
||||
sys.modules["folder_paths"] = types.SimpleNamespace(**attrs)
|
||||
|
||||
def _restore_folder_paths(self):
|
||||
if self._orig_folder_paths is None:
|
||||
sys.modules.pop("folder_paths", None)
|
||||
else:
|
||||
sys.modules["folder_paths"] = self._orig_folder_paths
|
||||
|
||||
def _restore_extra_env(self):
|
||||
if self._orig_extra is None:
|
||||
os.environ.pop(EXTRA_ROOTS_ENV, None)
|
||||
else:
|
||||
os.environ[EXTRA_ROOTS_ENV] = self._orig_extra
|
||||
|
||||
|
||||
class TestPublicPath(FolderPathsFixture):
|
||||
|
||||
def test_path_under_base_is_relative_posix(self):
|
||||
target = self.output_dir / "renders" / "a.png"
|
||||
self.assertEqual(_public_path(str(target)), "output/renders/a.png")
|
||||
|
||||
def test_base_itself_is_dot(self):
|
||||
self.assertEqual(_public_path(str(self.comfy_dir)), ".")
|
||||
|
||||
def test_path_outside_base_is_basename_only(self):
|
||||
secret = Path(self.tmpdir) / "secret" / "passwords.txt"
|
||||
secret.parent.mkdir()
|
||||
secret.write_text("x")
|
||||
self.assertEqual(_public_path(str(secret)), "passwords.txt")
|
||||
self.assertNotIn(self.tmpdir, _public_path(str(secret)))
|
||||
|
||||
def test_path_object_accepted(self):
|
||||
self.assertEqual(_public_path(self.output_dir), "output")
|
||||
|
||||
def test_falls_back_to_output_parent_without_base_path(self):
|
||||
self.install_folder_paths() # no base_path attribute
|
||||
self.assertEqual(_public_path(str(self.output_dir / "x.png")), "output/x.png")
|
||||
|
||||
def test_env_extra_root_keeps_its_own_name(self):
|
||||
extra = Path(self.tmpdir) / "nas" / "gallery"
|
||||
(extra / "sub").mkdir(parents=True)
|
||||
os.environ[EXTRA_ROOTS_ENV] = str(extra)
|
||||
self.assertEqual(_public_path(str(extra / "sub")), "gallery/sub")
|
||||
|
||||
def test_without_folder_paths_returns_basename(self):
|
||||
sys.modules.pop("folder_paths", None)
|
||||
self.assertEqual(_public_path(str(self.output_dir / "img.png")), "img.png")
|
||||
|
||||
def test_empty_or_none_gives_empty_string(self):
|
||||
self.assertEqual(_public_path(""), "")
|
||||
self.assertEqual(_public_path(None), "")
|
||||
|
||||
|
||||
class TestPublicError(FolderPathsFixture):
|
||||
|
||||
def test_oserror_uses_strerror_and_basename(self):
|
||||
err = FileNotFoundError(
|
||||
2, "No such file or directory", str(self.tmpdir) + "/a/b.png"
|
||||
)
|
||||
text = _public_error(err)
|
||||
self.assertIn("No such file or directory", text)
|
||||
self.assertIn("b.png", text)
|
||||
self.assertNotIn(self.tmpdir, text)
|
||||
|
||||
def test_oserror_without_filename(self):
|
||||
err = PermissionError(13, "Permission denied")
|
||||
self.assertEqual(_public_error(err), "Permission denied")
|
||||
|
||||
def test_oserror_without_strerror_uses_type_name(self):
|
||||
err = OSError()
|
||||
self.assertEqual(_public_error(err), "OSError")
|
||||
|
||||
def test_non_oserror_is_str(self):
|
||||
self.assertEqual(_public_error(ValueError("bad value")), "bad value")
|
||||
|
||||
def test_non_oserror_with_empty_message_uses_type_name(self):
|
||||
self.assertEqual(_public_error(RuntimeError()), "RuntimeError")
|
||||
|
||||
|
||||
# ── gzip middleware ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestGzipMiddleware(AioHTTPTestCase):
|
||||
|
||||
async def get_application(self):
|
||||
big = json.dumps({"data": "x" * 4000})
|
||||
|
||||
async def big_json(request):
|
||||
return web.json_response(text=big)
|
||||
|
||||
async def small_json(request):
|
||||
return web.json_response({"ok": True})
|
||||
|
||||
async def already_encoded(request):
|
||||
return web.Response(
|
||||
text=big, content_type="text/plain", headers={"Content-Encoding": "br"}
|
||||
)
|
||||
|
||||
async def binary(request):
|
||||
return web.Response(body=b"\x89PNG" * 1000, content_type="image/png")
|
||||
|
||||
async def incompressible(request):
|
||||
return web.Response(body=os.urandom(4096), content_type="text/plain")
|
||||
|
||||
async def streamed(request):
|
||||
resp = web.StreamResponse()
|
||||
await resp.prepare(request)
|
||||
await resp.write(b"x" * 4000)
|
||||
await resp.write_eof()
|
||||
return resp
|
||||
|
||||
app = web.Application(middlewares=[api_module._gzip_middleware])
|
||||
app.router.add_get("/prompt_manager/big", big_json)
|
||||
app.router.add_get("/prompt_manager/small", small_json)
|
||||
app.router.add_get("/prompt_manager/encoded", already_encoded)
|
||||
app.router.add_get("/prompt_manager/binary", binary)
|
||||
app.router.add_get("/prompt_manager/random", incompressible)
|
||||
app.router.add_get("/prompt_manager/stream", streamed)
|
||||
app.router.add_get("/other/big", big_json)
|
||||
return app
|
||||
|
||||
async def _get(self, path, accept="gzip"):
|
||||
headers = {"Accept-Encoding": accept} if accept else {}
|
||||
return await self.client.request(
|
||||
"GET",
|
||||
path,
|
||||
headers=headers,
|
||||
auto_decompress=False,
|
||||
skip_auto_headers=["Accept-Encoding"],
|
||||
)
|
||||
|
||||
async def test_large_json_is_gzipped(self):
|
||||
resp = await self._get("/prompt_manager/big")
|
||||
self.assertEqual(resp.headers.get("Content-Encoding"), "gzip")
|
||||
self.assertEqual(resp.headers.get("Vary"), "Accept-Encoding")
|
||||
import gzip
|
||||
|
||||
self.assertIn(b'"data"', gzip.decompress(await resp.read()))
|
||||
|
||||
async def test_small_body_not_compressed(self):
|
||||
resp = await self._get("/prompt_manager/small")
|
||||
self.assertIsNone(resp.headers.get("Content-Encoding"))
|
||||
|
||||
async def test_other_paths_untouched(self):
|
||||
resp = await self._get("/other/big")
|
||||
self.assertIsNone(resp.headers.get("Content-Encoding"))
|
||||
|
||||
async def test_without_accept_encoding(self):
|
||||
resp = await self._get("/prompt_manager/big", accept=None)
|
||||
self.assertIsNone(resp.headers.get("Content-Encoding"))
|
||||
|
||||
async def test_already_encoded_untouched(self):
|
||||
resp = await self._get("/prompt_manager/encoded")
|
||||
self.assertEqual(resp.headers.get("Content-Encoding"), "br")
|
||||
|
||||
async def test_binary_type_untouched(self):
|
||||
resp = await self._get("/prompt_manager/binary")
|
||||
self.assertIsNone(resp.headers.get("Content-Encoding"))
|
||||
|
||||
async def test_incompressible_body_untouched(self):
|
||||
resp = await self._get("/prompt_manager/random")
|
||||
self.assertIsNone(resp.headers.get("Content-Encoding"))
|
||||
|
||||
async def test_stream_response_untouched(self):
|
||||
resp = await self._get("/prompt_manager/stream")
|
||||
self.assertIsNone(resp.headers.get("Content-Encoding"))
|
||||
self.assertEqual(len(await resp.read()), 4000)
|
||||
|
||||
|
||||
# ── JSON body reader ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestReadJsonBody(AioHTTPTestCase):
|
||||
"""_read_json_body: chunked read with a byte cap, object-only payloads."""
|
||||
|
||||
async def get_application(self):
|
||||
async def echo(request):
|
||||
data, error = await api_module._read_json_body(request)
|
||||
if error is not None:
|
||||
return error
|
||||
return web.json_response({"success": True, "keys": sorted(data)})
|
||||
|
||||
async def small(request):
|
||||
data, error = await api_module._read_json_body(request, max_bytes=64)
|
||||
if error is not None:
|
||||
return error
|
||||
return web.json_response({"success": True, "data": data})
|
||||
|
||||
app = web.Application(client_max_size=4 * 1024 * 1024)
|
||||
app.router.add_post("/echo", echo)
|
||||
app.router.add_post("/small", small)
|
||||
return app
|
||||
|
||||
async def test_valid_object_is_parsed(self):
|
||||
resp = await self.client.request("POST", "/echo", json={"b": 1, "a": 2})
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual((await resp.json())["keys"], ["a", "b"])
|
||||
|
||||
async def test_default_cap_is_one_megabyte(self):
|
||||
self.assertEqual(api_module.JSON_BODY_MAX_BYTES, 1_000_000)
|
||||
body = json.dumps({"pad": "x" * 1_000_100}).encode()
|
||||
resp = await self.client.request(
|
||||
"POST", "/echo", data=body, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
self.assertEqual(resp.status, 413)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertIn("too large", data["error"].lower())
|
||||
|
||||
async def test_body_just_under_cap_is_accepted(self):
|
||||
payload = {"pad": "x" * (1_000_000 - 20)}
|
||||
body = json.dumps(payload).encode()
|
||||
self.assertLessEqual(len(body), 1_000_000)
|
||||
resp = await self.client.request(
|
||||
"POST", "/echo", data=body, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
|
||||
async def test_chunked_upload_without_content_length_is_capped(self):
|
||||
async def gen():
|
||||
for _ in range(10):
|
||||
yield b'{"pad": "' + b"x" * 20 + b'"}'
|
||||
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/small",
|
||||
data=gen(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
self.assertEqual(resp.status, 413)
|
||||
|
||||
async def test_declared_oversize_content_length_is_413(self):
|
||||
body = b'{"pad": "' + b"x" * 200 + b'"}'
|
||||
resp = await self.client.request(
|
||||
"POST", "/small", data=body, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
self.assertEqual(resp.status, 413)
|
||||
|
||||
async def test_invalid_json_is_400(self):
|
||||
for body in (b"{oops", b"\xff\xfe", b""):
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/echo",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
self.assertEqual(resp.status, 400, body)
|
||||
self.assertFalse((await resp.json())["success"])
|
||||
|
||||
async def test_non_object_json_is_400(self):
|
||||
for body in (b"[1, 2]", b'"text"', b"3", b"null", b"true"):
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/echo",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
self.assertEqual(resp.status, 400, body)
|
||||
self.assertIn("object", (await resp.json())["error"])
|
||||
|
||||
def test_helper_is_exposed_on_the_api_object(self):
|
||||
self.assertIs(PromptManagerAPI._read_json_body, api_module._read_json_body)
|
||||
|
||||
|
||||
class TestLongJobLock(unittest.TestCase):
|
||||
"""_acquire_long_job: one job per family at a time, released on exit."""
|
||||
|
||||
def setUp(self):
|
||||
self.api = PromptManagerAPI()
|
||||
|
||||
def _run(self, coro):
|
||||
import asyncio
|
||||
|
||||
return asyncio.run(coro)
|
||||
|
||||
def test_second_acquire_is_refused_until_first_releases(self):
|
||||
async def scenario():
|
||||
first = self.api._acquire_long_job("duplicates")
|
||||
self.assertIsNotNone(first)
|
||||
async with first:
|
||||
self.assertIsNone(self.api._acquire_long_job("duplicates"))
|
||||
# Other families are independent
|
||||
other = self.api._acquire_long_job("thumbnails")
|
||||
self.assertIsNotNone(other)
|
||||
async with other:
|
||||
pass
|
||||
again = self.api._acquire_long_job("duplicates")
|
||||
self.assertIsNotNone(again)
|
||||
async with again:
|
||||
pass
|
||||
|
||||
self._run(scenario())
|
||||
|
||||
def test_lock_is_released_when_the_job_raises(self):
|
||||
async def scenario():
|
||||
with self.assertRaises(RuntimeError):
|
||||
async with self.api._acquire_long_job("duplicates"):
|
||||
raise RuntimeError("boom")
|
||||
self.assertIsNotNone(self.api._acquire_long_job("duplicates"))
|
||||
|
||||
self._run(scenario())
|
||||
|
||||
def test_busy_response_is_409_with_envelope(self):
|
||||
resp = self.api._long_job_busy_response("Duplicate scan")
|
||||
self.assertEqual(resp.status, 409)
|
||||
data = json.loads(resp.body)
|
||||
self.assertFalse(data["success"])
|
||||
self.assertEqual(data["error"], "Duplicate scan already running")
|
||||
|
||||
|
||||
# ── UI and static routes ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class CoreRoutesTestCase(AioHTTPTestCase):
|
||||
|
||||
async def get_application(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
|
||||
async def tearDownAsync(self):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
self.api.db.close_all()
|
||||
shutil.rmtree(self.tmpdir, ignore_errors=True)
|
||||
|
||||
def _fake_root(self):
|
||||
"""Point _get_project_root at the temp dir and return its web/ dir."""
|
||||
web_dir = Path(self.tmpdir) / "web"
|
||||
(web_dir / "js").mkdir(parents=True)
|
||||
(web_dir / "lib").mkdir()
|
||||
orig_root = api_module._get_project_root
|
||||
api_module._get_project_root = lambda: self.tmpdir
|
||||
self.addCleanup(setattr, api_module, "_get_project_root", orig_root)
|
||||
return web_dir
|
||||
|
||||
|
||||
UI_PATHS = (
|
||||
"/prompt_manager/web",
|
||||
"/prompt_manager/gallery.html",
|
||||
"/prompt_manager/admin",
|
||||
"/prompt_manager/gallery",
|
||||
)
|
||||
|
||||
|
||||
class TestUiRoutes(CoreRoutesTestCase):
|
||||
|
||||
async def test_pages_are_served(self):
|
||||
for path in UI_PATHS:
|
||||
resp = await self.client.request("GET", path)
|
||||
self.assertEqual(resp.status, 200, path)
|
||||
self.assertIn("text/html", resp.headers["Content-Type"])
|
||||
self.assertIn("<", await resp.text())
|
||||
|
||||
async def test_pages_are_cached_after_first_read(self):
|
||||
await self.client.request("GET", "/prompt_manager/admin")
|
||||
self.assertTrue(any(p.endswith("admin.html") for p in self.api._html_cache))
|
||||
resp = await self.client.request("GET", "/prompt_manager/admin")
|
||||
self.assertEqual(resp.status, 200)
|
||||
|
||||
async def test_missing_pages_are_404(self):
|
||||
self._fake_root()
|
||||
for path in UI_PATHS:
|
||||
resp = await self.client.request("GET", path)
|
||||
self.assertEqual(resp.status, 404, path)
|
||||
|
||||
async def test_unreadable_pages_are_500(self):
|
||||
web_dir = self._fake_root()
|
||||
for name in ("index.html", "metadata.html", "admin.html", "gallery.html"):
|
||||
(web_dir / name).mkdir() # a directory cannot be read as a file
|
||||
for path in UI_PATHS:
|
||||
resp = await self.client.request("GET", path)
|
||||
self.assertEqual(resp.status, 500, path)
|
||||
self.assertNotIn(self.tmpdir, await resp.text())
|
||||
|
||||
async def test_test_route(self):
|
||||
resp = await self.client.request("GET", "/prompt_manager/test")
|
||||
self.assertTrue((await resp.json())["success"])
|
||||
|
||||
|
||||
class TestStaticRoutes(CoreRoutesTestCase):
|
||||
|
||||
async def test_js_served_with_mime(self):
|
||||
resp = await self.client.request("GET", "/prompt_manager/js/admin.js")
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertIn("application/javascript", resp.headers["Content-Type"])
|
||||
|
||||
async def test_lib_served(self):
|
||||
web_dir = self._fake_root()
|
||||
(web_dir / "lib" / "x.css").write_text("body{}")
|
||||
resp = await self.client.request("GET", "/prompt_manager/lib/x.css")
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertIn("text/css", resp.headers["Content-Type"])
|
||||
self.assertEqual(await resp.text(), "body{}")
|
||||
|
||||
async def test_real_lib_file_served_with_content_type(self):
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/lib/viewerjs/viewer.min.css"
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertIn("text/css", resp.headers["Content-Type"])
|
||||
|
||||
async def test_allowed_asset_types_get_their_mime(self):
|
||||
web_dir = self._fake_root()
|
||||
expected = {
|
||||
"a.map": "application/json",
|
||||
"f.woff": "font/woff",
|
||||
"f.woff2": "font/woff2",
|
||||
"f.ttf": "font/ttf",
|
||||
"i.png": "image/png",
|
||||
"i.svg": "image/svg+xml",
|
||||
"i.ico": "image/x-icon",
|
||||
"p.html": "text/html",
|
||||
}
|
||||
for name in expected:
|
||||
(web_dir / "lib" / name).write_bytes(b"\x00\x01")
|
||||
for name, mime in expected.items():
|
||||
resp = await self.client.request("GET", f"/prompt_manager/lib/{name}")
|
||||
self.assertEqual(resp.status, 200, name)
|
||||
self.assertIn(mime, resp.headers["Content-Type"], name)
|
||||
|
||||
async def test_disallowed_extensions_are_403_even_when_present(self):
|
||||
web_dir = self._fake_root()
|
||||
for name in ("x.bin", "x.py", "x.json", "noext", "x.JS.bak"):
|
||||
(web_dir / "lib" / name).write_bytes(b"\x00")
|
||||
(web_dir / "js" / name).write_bytes(b"\x00")
|
||||
for prefix in ("lib", "js"):
|
||||
for name in ("x.bin", "x.py", "x.json", "noext", "x.JS.bak"):
|
||||
resp = await self.client.request(
|
||||
"GET", f"/prompt_manager/{prefix}/{name}"
|
||||
)
|
||||
self.assertEqual(resp.status, 403, f"{prefix}/{name}")
|
||||
|
||||
async def test_traversal_is_403(self):
|
||||
for prefix in ("lib", "js"):
|
||||
for encoded in (
|
||||
"..%2F..%2Fpyproject.toml",
|
||||
"..%5C..%5Cpyproject.toml",
|
||||
"sub%2F..%2F..%2F..%2Fpyproject.toml",
|
||||
"%2E%2E%2Fpyproject.toml",
|
||||
):
|
||||
resp = await self.client.request(
|
||||
"GET", f"/prompt_manager/{prefix}/{encoded}"
|
||||
)
|
||||
self.assertEqual(resp.status, 403, f"{prefix}/{encoded}")
|
||||
|
||||
async def test_absolute_drive_and_unc_paths_are_rejected_on_every_os(self):
|
||||
# Built literally so the Windows-only escapes are exercised everywhere:
|
||||
# os.path.join(root, "C:/x") drops root on Windows, "\\x" and UNC
|
||||
# shares re-anchor the path as well.
|
||||
candidates = (
|
||||
"C:/Windows/win.ini",
|
||||
"C:%5CWindows%5Cwin.ini",
|
||||
"c:win.ini",
|
||||
"%5C%5Cserver%5Cshare%5Cx.js",
|
||||
"%2F%2Fserver%2Fshare%2Fx.js",
|
||||
"%5CWindows%5Cwin.ini",
|
||||
"%2Fetc%2Fpasswd",
|
||||
)
|
||||
for prefix in ("lib", "js"):
|
||||
for encoded in candidates:
|
||||
resp = await self.client.request(
|
||||
"GET", f"/prompt_manager/{prefix}/{encoded}"
|
||||
)
|
||||
self.assertIn(resp.status, (403, 404), f"{prefix}/{encoded}")
|
||||
self.assertNotIn("[extensions]", await resp.text())
|
||||
|
||||
async def test_symlink_escaping_web_dir_is_403(self):
|
||||
web_dir = self._fake_root()
|
||||
secret = Path(self.tmpdir) / "secret.js"
|
||||
secret.write_text("secret")
|
||||
link = web_dir / "lib" / "link.js"
|
||||
try:
|
||||
os.symlink(secret, link)
|
||||
except (OSError, NotImplementedError):
|
||||
self.skipTest("symlinks not available")
|
||||
resp = await self.client.request("GET", "/prompt_manager/lib/link.js")
|
||||
self.assertEqual(resp.status, 403)
|
||||
|
||||
async def test_empty_path_is_not_served(self):
|
||||
for prefix in ("lib", "js"):
|
||||
resp = await self.client.request("GET", f"/prompt_manager/{prefix}/")
|
||||
self.assertIn(resp.status, (403, 404), prefix)
|
||||
|
||||
async def test_missing_is_404(self):
|
||||
for prefix in ("lib", "js"):
|
||||
resp = await self.client.request("GET", f"/prompt_manager/{prefix}/nope.js")
|
||||
self.assertEqual(resp.status, 404, prefix)
|
||||
|
||||
async def test_directory_is_404(self):
|
||||
web_dir = self._fake_root()
|
||||
(web_dir / "lib" / "dir.js").mkdir()
|
||||
resp = await self.client.request("GET", "/prompt_manager/lib/dir.js")
|
||||
self.assertEqual(resp.status, 404)
|
||||
|
||||
async def test_root_of_static_dir_is_not_listed(self):
|
||||
resp = await self.client.request("GET", "/prompt_manager/lib/tailwind")
|
||||
self.assertEqual(resp.status, 403)
|
||||
|
||||
|
||||
# ── instance helpers ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestInstanceHelpers(FolderPathsFixture):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
self.addCleanup(self.api.db.close_all) # runs before the rmtree cleanup
|
||||
orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
self.addCleanup(setattr, GalleryConfig, "MONITORING_DIRECTORIES", orig_dirs)
|
||||
|
||||
def _run(self, coro):
|
||||
import asyncio
|
||||
|
||||
return asyncio.run(coro)
|
||||
|
||||
def test_run_in_executor_with_kwargs(self):
|
||||
def add(a, b=0):
|
||||
return a + b
|
||||
|
||||
self.assertEqual(self._run(self.api._run_in_executor(add, 1, b=2)), 3)
|
||||
self.assertEqual(self._run(self.api._run_in_executor(add, 1, 2)), 3)
|
||||
|
||||
def test_invalidate_gallery_cache(self):
|
||||
self.api._gallery_cache = {"x": 1}
|
||||
self.api.invalidate_gallery_cache()
|
||||
self.assertEqual(self.api._gallery_cache, {})
|
||||
|
||||
def test_enrich_images_adds_urls_and_thumbnails(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
img = self.output_dir / "sub" / "a.png"
|
||||
img.parent.mkdir()
|
||||
img.write_bytes(b"x")
|
||||
thumb = self.output_dir / "thumbnails" / "sub" / "a_thumb.png"
|
||||
thumb.parent.mkdir(parents=True)
|
||||
thumb.write_bytes(b"t")
|
||||
outside = Path(self.tmpdir) / "elsewhere.png"
|
||||
outside.write_bytes(b"o")
|
||||
|
||||
images = [
|
||||
{"id": 7, "image_path": str(img)},
|
||||
{"id": 8, "image_path": str(outside)},
|
||||
{"id": 9, "image_path": ""},
|
||||
{"image_path": str(self.output_dir / "missing.png")},
|
||||
]
|
||||
result = self.api._enrich_images(images)
|
||||
|
||||
self.assertEqual(result[0]["url"], "/prompt_manager/images/serve/sub/a.png")
|
||||
self.assertRegex(
|
||||
result[0]["thumbnail_url"],
|
||||
r"^/prompt_manager/images/serve/thumbnails/sub/a_thumb\.png\?v=\d+$",
|
||||
)
|
||||
self.assertEqual(result[0]["relative_path"], os.path.join("sub", "a.png"))
|
||||
self.assertEqual(result[1]["url"], "/prompt_manager/images/8/file")
|
||||
self.assertNotIn("relative_path", result[1])
|
||||
self.assertNotIn("url", result[2])
|
||||
self.assertEqual(result[3]["url"], "/prompt_manager/images/serve/missing.png")
|
||||
self.assertNotIn("thumbnail_url", result[3])
|
||||
|
||||
def test_enrich_prompt_images_walks_prompts(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
prompts = [
|
||||
{"images": [{"id": 1, "image_path": str(self.output_dir / "p.png")}]},
|
||||
{},
|
||||
]
|
||||
result = self.api._enrich_prompt_images(prompts)
|
||||
self.assertEqual(
|
||||
result[0]["images"][0]["url"], "/prompt_manager/images/serve/p.png"
|
||||
)
|
||||
|
||||
def test_clean_nan_recursive(self):
|
||||
nan = float("nan")
|
||||
cleaned = self.api._clean_nan_recursive(
|
||||
{"a": nan, "b": [nan, 1.5, {"c": nan}], "d": "x"}
|
||||
)
|
||||
self.assertEqual(cleaned, {"a": None, "b": [None, 1.5, {"c": None}], "d": "x"})
|
||||
|
||||
def test_find_output_dir_uses_configured_and_caches(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
self.api._cached_output_dir = None
|
||||
first = self.api._find_comfyui_output_dir()
|
||||
self.assertEqual(first, str(self.output_dir.resolve()))
|
||||
GalleryConfig.MONITORING_DIRECTORIES = []
|
||||
self.assertEqual(self.api._find_comfyui_output_dir(), first)
|
||||
|
||||
def test_find_output_dir_skips_configured_root_outside_comfyui(self):
|
||||
outside = Path(self.tmpdir) / "outside"
|
||||
outside.mkdir()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
|
||||
self.api._cached_output_dir = None
|
||||
result = self.api._find_comfyui_output_dir()
|
||||
self.assertNotEqual(
|
||||
os.path.normcase(str(result or "")), os.path.normcase(str(outside))
|
||||
)
|
||||
|
||||
def test_find_output_dir_skips_missing_configured(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [os.path.join(self.tmpdir, "nope")]
|
||||
self.api._cached_output_dir = None
|
||||
result = self.api._find_comfyui_output_dir()
|
||||
self.assertNotEqual(result, os.path.join(self.tmpdir, "nope"))
|
||||
|
||||
|
||||
# ── metadata parsing ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestMetadataParsing(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.api = PromptManagerAPI()
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, self.tmpdir, True)
|
||||
|
||||
def test_extract_metadata_from_png_and_garbage(self):
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
path = os.path.join(self.tmpdir, "m.png")
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("prompt", "{}")
|
||||
Image.new("RGB", (1, 1)).save(path, pnginfo=info)
|
||||
self.assertEqual(self.api._extract_comfyui_metadata(path), {"prompt": "{}"})
|
||||
|
||||
garbage = os.path.join(self.tmpdir, "g.png")
|
||||
with open(garbage, "wb") as f:
|
||||
f.write(b"nope")
|
||||
self.assertEqual(self.api._extract_comfyui_metadata(garbage), {})
|
||||
|
||||
def test_parse_a1111_parameters(self):
|
||||
meta = {"parameters": "a cat\nNegative prompt: dog\nSteps: 20"}
|
||||
parsed = self.api._parse_comfyui_prompt(meta)
|
||||
self.assertEqual(parsed["positive_prompt"], "a cat")
|
||||
self.assertEqual(parsed["negative_prompt"], "dog")
|
||||
self.assertEqual(parsed["parameters"]["parameters"], meta["parameters"])
|
||||
self.assertEqual(self.api._extract_readable_prompt(parsed), "a cat")
|
||||
|
||||
def test_parse_comfyui_fields(self):
|
||||
meta = {
|
||||
"prompt": json.dumps({"1": {"class_type": "X", "inputs": {}}}),
|
||||
"workflow": "not json",
|
||||
"steps": "20",
|
||||
"cfg": "not-json-either",
|
||||
}
|
||||
parsed = self.api._parse_comfyui_prompt(meta)
|
||||
self.assertIsInstance(parsed["prompt"], dict)
|
||||
self.assertEqual(parsed["workflow"], "not json")
|
||||
self.assertEqual(parsed["parameters"]["steps"], 20)
|
||||
self.assertEqual(parsed["parameters"]["cfg"], "not-json-either")
|
||||
|
||||
parsed = self.api._parse_comfyui_prompt(
|
||||
{"prompt": "plain text", "workflow": json.dumps({"nodes": []})}
|
||||
)
|
||||
self.assertEqual(parsed["prompt"], "plain text")
|
||||
self.assertEqual(parsed["workflow"], {"nodes": []})
|
||||
|
||||
def test_extract_readable_prompt_variants(self):
|
||||
extract = self.api._extract_readable_prompt
|
||||
self.assertEqual(extract({"prompt": "plain"}), "plain")
|
||||
self.assertEqual(extract({"prompt": ["a", "", "b"]}), "a b")
|
||||
self.assertEqual(extract({"prompt": 42}), "42")
|
||||
self.assertEqual(
|
||||
extract({"prompt": None, "parameters": {"positive": ["x", "y"]}}), "x y"
|
||||
)
|
||||
self.assertIsNone(extract({"prompt": {}, "parameters": {}}))
|
||||
|
||||
graph = {
|
||||
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "neg"}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
||||
},
|
||||
}
|
||||
self.assertEqual(extract({"prompt": graph}), "pos")
|
||||
self.assertEqual(extract({"prompt": {"x": 1}, "workflow": graph}), "pos")
|
||||
|
||||
def test_node_inputs_and_text_lookup(self):
|
||||
self.assertEqual(self.api._get_node_inputs("nope"), {})
|
||||
self.assertEqual(self.api._get_node_inputs({"inputs": 5}), {})
|
||||
self.assertEqual(self.api._get_node_inputs({"inputs": {"a": 1}}), {"a": 1})
|
||||
listed = self.api._get_node_inputs(
|
||||
{"inputs": [{"name": "text", "link": None}, {"nope": 1}, "junk"]}
|
||||
)
|
||||
self.assertEqual(list(listed), ["text"])
|
||||
|
||||
self.assertIsNone(self.api._find_text_in_node(None))
|
||||
self.assertEqual(self.api._find_text_in_node({"inputs": {"text": "hi"}}), "hi")
|
||||
self.assertEqual(
|
||||
self.api._find_text_in_node(
|
||||
{"type": "CLIPTextEncode", "widgets_values": ["w"]}
|
||||
),
|
||||
"w",
|
||||
)
|
||||
self.assertIsNone(
|
||||
self.api._find_text_in_node(
|
||||
{"type": "CLIPTextEncode", "widgets_values": [" "]}
|
||||
)
|
||||
)
|
||||
self.assertIsNone(self.api._find_text_in_node({"type": "KSampler"}))
|
||||
|
||||
def test_positive_prompt_from_nodes_array_and_fallbacks(self):
|
||||
extract = self.api._extract_positive_prompt_from_comfyui_data
|
||||
self.assertIsNone(extract("x"))
|
||||
self.assertIsNone(extract({}))
|
||||
self.assertIsNone(extract({"nodes": []}))
|
||||
|
||||
workflow = {
|
||||
"nodes": [
|
||||
{"id": 1, "type": "CLIPTextEncode", "widgets_values": ["pos"]},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "CLIPTextEncode",
|
||||
"widgets_values": ["neg"],
|
||||
"title": "Negative",
|
||||
},
|
||||
{"id": 3, "type": "KSampler", "inputs": [{"name": "seed", "link": 4}]},
|
||||
"junk",
|
||||
]
|
||||
}
|
||||
self.assertEqual(extract(workflow), "pos")
|
||||
|
||||
only_negative = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": 2,
|
||||
"type": "CLIPTextEncode",
|
||||
"widgets_values": ["neg"],
|
||||
"title": "neg",
|
||||
},
|
||||
]
|
||||
}
|
||||
self.assertEqual(extract(only_negative), "neg")
|
||||
|
||||
bad_link = {
|
||||
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": "oops", "negative": []},
|
||||
},
|
||||
"4": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["zz", 0], "negative": []},
|
||||
},
|
||||
}
|
||||
self.assertEqual(extract(bad_link), "pos")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||