Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a8b5a5aee2 | ||
|
|
2696781d6b | ||
|
|
7517f1fe80 | ||
|
|
8a5638221f | ||
|
|
0b2b274583 | ||
|
|
eacd239d23 | ||
|
|
04ade0dae5 | ||
|
|
3c912c3a1d | ||
|
|
0eefe12131 | ||
|
|
b2b6f6badb | ||
|
|
00cfd979bf | ||
|
|
53dfbe4999 | ||
|
|
88d1ac3666 | ||
|
|
0d5f76bf89 | ||
|
|
8c9fcf5225 | ||
|
|
69fe4ceaf1 | ||
|
|
b21497f664 | ||
|
|
29a8d9286f | ||
|
|
d3ccc24055 | ||
|
|
8756d720ae | ||
|
|
33e9201ac8 | ||
|
|
3f2a1d5421 | ||
|
|
49f8cfbde6 | ||
|
|
913fc5354d | ||
|
|
fb8e8ec8d5 | ||
|
|
59b41a19a4 | ||
|
|
aaf49230b2 | ||
|
|
b37d90882b | ||
|
|
71608332ab | ||
|
|
48e452dd67 | ||
|
|
377b95681c | ||
|
|
b4f388f378 | ||
|
|
de815c5c63 | ||
|
|
88e73b81c3 | ||
|
|
4145f247b7 | ||
|
|
1e3999e32f | ||
|
|
5cdd808bd8 | ||
|
|
aa63fffb36 | ||
|
|
8727eab44e | ||
|
|
19c6084216 | ||
|
|
63747d9c9c | ||
|
|
f7ec8e5073 | ||
|
|
0e635cd37f | ||
|
|
e596c6de5a | ||
|
|
387b5ceb0b | ||
|
|
5369980d57 | ||
|
|
61011eea96 | ||
|
|
d706f0eb00 | ||
|
|
77637e0626 | ||
|
|
9c9917961c | ||
|
|
a32bb7f913 | ||
|
|
4c6493fe63 | ||
|
|
c12beff2d2 | ||
|
|
5e2872d5f8 | ||
|
|
5a8f23ea2b | ||
|
|
12492185d2 | ||
|
|
c99af98e09 | ||
|
|
a8a8946edc | ||
|
|
16d2ed916d | ||
|
|
31047dc9da | ||
|
|
e15f716563 | ||
|
|
a6fbfe06e0 | ||
|
|
11a8d9d218 | ||
|
|
ebc9903bf2 | ||
|
|
7c57ccaae4 | ||
|
|
3653d2883f | ||
|
|
9c5870f938 | ||
|
|
005efdd7c5 | ||
|
|
78682dd0cb | ||
|
|
29198a5fd7 | ||
|
|
5c09936c3a | ||
|
|
ba35d58e81 | ||
|
|
ef95266cf4 | ||
|
|
b4228dc340 | ||
|
|
88f1485cc9 | ||
|
|
16ef5bf5db | ||
|
|
983862eec3 | ||
|
|
acc1570629 | ||
|
|
30c0547009 | ||
|
|
13e24c3a57 | ||
|
|
cd00cce1f8 | ||
|
|
a26d060d5f | ||
|
|
e9594ec428 | ||
|
|
b6b2356389 | ||
|
|
c3375c61af | ||
|
|
d39c0c89fe | ||
|
|
e1a0187f8f | ||
|
|
ac181210d1 | ||
|
|
e47ce1dd43 | ||
|
|
1c9f9568b5 | ||
|
|
29d9790c8a | ||
|
|
15c6f13ac3 | ||
|
|
b62511a5b3 | ||
|
|
3c7c1ca267 | ||
|
|
0e456783e3 | ||
|
|
b93d7ffe18 | ||
|
|
026107923b | ||
|
|
3bb42c3d2f | ||
|
|
5577efa485 | ||
|
|
121802da9e | ||
|
|
a40f09c7d9 | ||
|
|
95421b410f | ||
|
|
a3a1bb2446 | ||
|
|
3e9b4e1d6e | ||
|
|
dcadd6d605 | ||
|
|
720021b940 | ||
|
|
65635e3e95 | ||
|
|
024cb92f1f | ||
|
|
b2cfa6b8d9 | ||
|
|
babfecfaf1 | ||
|
|
66ab2f9606 | ||
|
|
3af42597d0 | ||
|
|
f0942b25b3 | ||
|
|
b82da68c8a | ||
|
|
ac2c097bfa | ||
|
|
c61537509f | ||
|
|
402c26db7f | ||
|
|
995d508979 | ||
|
|
a9ee9239f3 | ||
|
|
06ddb6540b | ||
|
|
e6ab7910a7 | ||
|
|
3473644b75 | ||
|
|
a98b8ebb18 | ||
|
|
ff969b8455 | ||
|
|
bebb9f6e7f | ||
|
|
5e3d357304 | ||
|
|
49a2e9d8fd | ||
|
|
cc44174547 | ||
|
|
7e950bae86 | ||
|
|
9e011874ea | ||
|
|
5225d8ea17 | ||
|
|
d92d115e96 | ||
|
|
59d7d1efa8 | ||
|
|
f8411cad14 | ||
|
|
2e7674a387 | ||
|
|
88c14feea7 | ||
|
|
f4fddec96d | ||
|
|
92b06c3f41 | ||
|
|
3e910d8fd9 | ||
|
|
e149843077 | ||
|
|
ccff1bd256 | ||
|
|
b6f63ab7f8 | ||
|
|
1ce384d2d8 | ||
|
|
a603b8175d | ||
|
|
7782c50b5c | ||
|
|
c9f55c28c1 | ||
|
|
fac1de762a | ||
|
|
6bbe02eecd | ||
|
|
18d1d80b15 | ||
|
|
c5aeb01db0 | ||
|
|
919d7a1f42 | ||
|
|
a185d5091f | ||
|
|
b96dab6eb5 | ||
|
|
2650192596 | ||
|
|
1fade6be5d | ||
|
|
d319916fb9 | ||
|
|
628b074d84 | ||
|
|
8757cbadb6 | ||
|
|
fa8cd0c77b | ||
|
|
a3f5ae33f2 | ||
|
|
ebf5b3e61a | ||
|
|
4480a17488 | ||
|
|
0b1af84172 | ||
|
|
ee63ca5ab5 | ||
|
|
8a3d4d2e48 | ||
|
|
2d91535777 | ||
|
|
bd9faa4ee6 | ||
|
|
08a6563378 | ||
|
|
128f37ac48 | ||
|
|
abc2fd6ee5 | ||
|
|
b99d14916e | ||
|
|
40e6079c34 | ||
|
|
a020c826ad | ||
|
|
bc62af2c35 | ||
|
|
7107740f71 | ||
|
|
64af37a73e | ||
|
|
9a096e73eb | ||
|
|
7bf9363777 | ||
|
|
445c4500b5 | ||
|
|
af3606677f | ||
|
|
5b52d88daa | ||
|
|
455b122e05 | ||
|
|
5936cda2e7 | ||
|
|
e9ca2935e6 | ||
|
|
de1d23911b | ||
|
|
4569fbeb98 | ||
|
|
7601b266ca | ||
|
|
71f5edb5bb | ||
|
|
102b0882e6 | ||
|
|
58d4925c28 | ||
|
|
22a5d248c9 | ||
|
|
70ffb2a661 | ||
|
|
9ac769fa22 | ||
|
|
bd6bbb7b6f | ||
|
|
f0ed280d6d | ||
|
|
47a9caa8bd | ||
|
|
38fb60ac78 | ||
|
|
44a07af659 | ||
|
|
7259978dc8 | ||
|
|
7260d2d04a | ||
|
|
748b3e1636 | ||
|
|
0b88844001 | ||
|
|
95ff2ed97a | ||
|
|
0eef846640 | ||
|
|
3476a41c24 | ||
|
|
2c057ca539 | ||
|
|
6754b429f3 | ||
|
|
86c3a9aba7 | ||
|
|
96bdc8bb38 | ||
|
|
083f20a36f | ||
|
|
09273849c5 | ||
|
|
f47619b37a | ||
|
|
46e5ffb69c | ||
|
|
b62a643920 | ||
|
|
6ee057e3d8 | ||
|
|
782e0709e2 | ||
|
|
4692503594 | ||
|
|
dcab859972 | ||
|
|
056604ebdb | ||
|
|
dd7c8e79eb | ||
|
|
9d2c335f2b | ||
|
|
a0a404e676 | ||
|
|
0592906b1c |
@@ -0,0 +1,24 @@
|
||||
# Files and folders excluded from the ComfyUI Registry package archive.
|
||||
# Everything listed here stays in the GitHub repository but is not bundled
|
||||
# when publishing via .github/workflows/publish_action.yml.
|
||||
|
||||
# Development, tests & CI
|
||||
.github/
|
||||
tests/
|
||||
|
||||
# Contributor & agent documentation (not needed at runtime)
|
||||
AGENTS.md
|
||||
CLAUDE.md
|
||||
GEMINI.md
|
||||
CONTRIBUTING.md
|
||||
STATUS.md
|
||||
LORA_LOADER_AUDIT.md
|
||||
|
||||
# Environment variables & local credentials
|
||||
.env
|
||||
.env.*
|
||||
*.env
|
||||
|
||||
# Python bytecode caches
|
||||
__pycache__/
|
||||
*.pyc
|
||||
+32
-3
@@ -1,8 +1,37 @@
|
||||
# ComfyUI-FlowControl Environment Configuration Example
|
||||
# Copy or rename this file to .env to configure custom API keys.
|
||||
# ComfyUI-SaturnNodes Environment Configuration Example
|
||||
# Copy or rename this file to .env to configure custom API keys and default behaviors.
|
||||
# Runtime location: ComfyUI/user/default/SaturnNodes/.env
|
||||
|
||||
# Optional Civitai API Key for auto-downloading LoRA preview thumbnails
|
||||
CIVITAI_API_KEY=
|
||||
|
||||
# Optional TMDB API Key for celebrity image preview search
|
||||
# Optional TMDB API Key / Read Access Token for actor / media preview image search
|
||||
TMDB_API_KEY=
|
||||
|
||||
# Feature toggles
|
||||
ENABLE_CIVITAI_SCRAPING=false
|
||||
ENABLE_TMDB_SCRAPING=false
|
||||
ENABLE_LORA_USAGE=true
|
||||
ENABLE_PERSISTENT_QUEUE=true
|
||||
ENABLE_TRAY_ICON=false
|
||||
ENABLE_ASSETS_RESTORE=true
|
||||
RESTORE_ASSETS_COUNT=64
|
||||
RESTORE_ASSETS_RECURSIVE=true
|
||||
SHOW_RESTORED_GENERATION_TIME=true
|
||||
|
||||
# Visual image loader metadata scanning (0 disables prompt extraction)
|
||||
IMAGE_LOADER_METADATA_LIMIT=128
|
||||
|
||||
# Launch defaults
|
||||
DEFAULT_PAUSE_STATE=Running
|
||||
DEFAULT_PAUSE_MODE=after_finish
|
||||
PERSISTENT_QUEUE_RESTORED_STATE=Match Default
|
||||
PERSISTENT_QUEUE_MAX_RETRIES=0
|
||||
PERSIST_PAUSE_STATE=false
|
||||
|
||||
# Prompt iterator privacy
|
||||
CLEAR_PROMPT_ITERATOR_ON_LAUNCH=false
|
||||
|
||||
# Security gates (opt-in, both disabled by default)
|
||||
ALLOW_PROCESS_MANAGEMENT=false
|
||||
ENABLE_LOCAL_FILE_EXECUTION=false
|
||||
|
||||
@@ -1,29 +1,30 @@
|
||||
---
|
||||
name: Bug report
|
||||
about: Create a report to help us improve ComfyUI-FlowControl
|
||||
name: 🐛 Bug Report
|
||||
about: Create a report to help us fix a bug or unexpected behavior
|
||||
title: '[BUG] '
|
||||
labels: 'bug'
|
||||
labels: bug
|
||||
assignees: ''
|
||||
---
|
||||
|
||||
**Describe the bug**
|
||||
**Describe the Bug**
|
||||
A clear and concise description of what the bug is.
|
||||
|
||||
**ComfyUI & System Information**
|
||||
- ComfyUI Version:
|
||||
- Frontend Version (V1 / V2):
|
||||
- OS (Windows / Linux / macOS):
|
||||
- Browser:
|
||||
|
||||
**To Reproduce**
|
||||
Steps to reproduce the behavior:
|
||||
1. Go to '...'
|
||||
2. Click on '....'
|
||||
3. Scroll down to '....'
|
||||
1. Load node '...'
|
||||
2. Connect input '...'
|
||||
3. Queue prompt
|
||||
4. See error
|
||||
|
||||
**Expected behavior**
|
||||
A clear and concise description of what you expected to happen.
|
||||
**Expected Behavior**
|
||||
A clear description of what you expected to happen.
|
||||
|
||||
**Console Logs / Screenshots**
|
||||
If applicable, paste terminal console logs or screenshots to help explain your problem.
|
||||
**Environment (please complete the following information):**
|
||||
- ComfyUI Version / UI Version: [e.g. Nodes 2.0 Vue UI or LiteGraph V1]
|
||||
- OS: [e.g. Windows 11, Ubuntu 22.04, macOS Sonoma]
|
||||
- Browser: [e.g. Chrome, Firefox, Edge]
|
||||
|
||||
**Console Logs / Traceback**
|
||||
```text
|
||||
Paste full terminal error traceback here (if applicable)
|
||||
```
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: Feature request
|
||||
about: Suggest an idea for ComfyUI-FlowControl
|
||||
about: Suggest an idea for ComfyUI-SaturnNodes
|
||||
title: '[FEAT] '
|
||||
labels: 'enhancement'
|
||||
assignees: ''
|
||||
|
||||
@@ -2,9 +2,9 @@ name: Lint and Syntax Check
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ main ]
|
||||
branches: [ "**" ]
|
||||
pull_request:
|
||||
branches: [ main ]
|
||||
branches: [ "**" ]
|
||||
|
||||
jobs:
|
||||
lint:
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
name: Publish to ComfyUI Registry
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
if: github.ref == 'refs/heads/main'
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
- name: Verify registry token
|
||||
env:
|
||||
REGISTRY_ACCESS_TOKEN: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
run: |
|
||||
if [ -z "$REGISTRY_ACCESS_TOKEN" ]; then
|
||||
echo "::error::REGISTRY_ACCESS_TOKEN is not configured. Create an API key at https://registry.comfy.org/nodes for the publisher 'kofiblto' and store it under Settings -> Secrets and variables -> Actions."
|
||||
exit 1
|
||||
fi
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
@@ -8,11 +8,27 @@ on:
|
||||
jobs:
|
||||
release:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Extract release notes from CHANGELOG.md
|
||||
run: |
|
||||
python - <<'PY'
|
||||
import os, re, pathlib
|
||||
tag = os.environ["GITHUB_REF_NAME"].lstrip("v")
|
||||
changelog = pathlib.Path("CHANGELOG.md").read_text(encoding="utf-8")
|
||||
pattern = re.compile(rf"^## \[{re.escape(tag)}\][^\n]*\n(.*?)(?=^## \[|\Z)", re.M | re.S)
|
||||
match = pattern.search(changelog)
|
||||
body = match.group(1).strip() if match else ""
|
||||
if not body:
|
||||
body = f"See CHANGELOG.md for the full list of changes in v{tag}."
|
||||
pathlib.Path("release_notes.md").write_text(body + "\n", encoding="utf-8")
|
||||
PY
|
||||
- name: Create GitHub Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
generate_release_notes: true
|
||||
body_path: release_notes.md
|
||||
generate_release_notes: false
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
name: Run Unit Tests
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ "**" ]
|
||||
pull_request:
|
||||
branches: [ "**" ]
|
||||
|
||||
jobs:
|
||||
test:
|
||||
name: Test on Python ${{ matrix.python-version }} (${{ matrix.os }})
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest]
|
||||
python-version: ['3.10', '3.11', '3.12']
|
||||
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
- name: Install Python Dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install Pillow torch numpy aiohttp piexif
|
||||
|
||||
- name: Execute SaturnNodes Unit Test Suite
|
||||
run: |
|
||||
python tests/run_tests.py
|
||||
@@ -13,9 +13,18 @@ persistent_queue.json.tmp
|
||||
lora_usage.json
|
||||
failed_scrapes.json
|
||||
image_prompts_cache.json
|
||||
assets_restore_debug.json
|
||||
assets_restore_debug.log
|
||||
user/
|
||||
*.json.tmp
|
||||
|
||||
# IDE and OS metadata
|
||||
.idea/
|
||||
.vscode/
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Downloaded models and checkpoints
|
||||
models/
|
||||
*.onnx
|
||||
*.onnx.tmp
|
||||
|
||||
@@ -0,0 +1,295 @@
|
||||
# 🤖 Agent & Developer Standard Operating Procedure (SOP)
|
||||
## `ComfyUI-SaturnNodes` Contributor & Automation Guide
|
||||
|
||||
This guide defines the engineering checklist and architecture standards for all AI agents (Antigravity, Claude, Gemini, Cursor, Copilot) and human developers contributing to **`ComfyUI-SaturnNodes`**.
|
||||
|
||||
For Git workflow, branch protection, and Pull Request guidelines, see **[`CONTRIBUTING.md`](./CONTRIBUTING.md)**.
|
||||
|
||||
---
|
||||
|
||||
## 📚 Repository Markdown Index
|
||||
|
||||
| File | Purpose & Contents | Link |
|
||||
| :--- | :--- | :--- |
|
||||
| **`AGENTS.md`** | Authoritative developer & agent standard operating procedure, engineering checklists, architecture rules. | [AGENTS.md](./AGENTS.md) |
|
||||
| **`CONTRIBUTING.md`** | Community contribution guidelines, strict branch rules, PR submission workflow. | [CONTRIBUTING.md](./CONTRIBUTING.md) |
|
||||
| **`README.md`** | Main project homepage, installation guide, complete node catalog with parameter tables and UI settings guide. | [README.md](./README.md) |
|
||||
| **`WALKTHROUGH.md`** | Step-by-step user tutorials, example workflows, visual load guides, and batch prompting patterns. | [WALKTHROUGH.md](./WALKTHROUGH.md) |
|
||||
| **`CHANGELOG.md`** | Strict Keep-A-Changelog semantic version release history across all updates. | [CHANGELOG.md](./CHANGELOG.md) |
|
||||
| **`SECURITY.md`** | Security policies, supported versions, and private vulnerability disclosure instructions. | [SECURITY.md](./SECURITY.md) |
|
||||
| **`SECURITY_AUDIT_PLAYBOOK.md`** | Authoritative ComfyUI custom node security rulebook, PR #3200 review precedents, threat model, and coding standards. | [SECURITY_AUDIT_PLAYBOOK.md](./SECURITY_AUDIT_PLAYBOOK.md) |
|
||||
| **`CLAUDE.md`** | AI agent instruction pointer. | [CLAUDE.md](./CLAUDE.md) |
|
||||
| **`GEMINI.md`** | AI agent instruction pointer. | [GEMINI.md](./GEMINI.md) |
|
||||
| **`bug_report.md`** | GitHub issue template for reporting reproducible defects. | [bug_report.md](./.github/ISSUE_TEMPLATE/bug_report.md) |
|
||||
| **`feature_request.md`** | GitHub issue template for proposing new nodes or UI features. | [feature_request.md](./.github/ISSUE_TEMPLATE/feature_request.md) |
|
||||
|
||||
---
|
||||
|
||||
## 📑 Table of Contents
|
||||
1. [Architecture & Directory Layout](#1-architecture--directory-layout)
|
||||
2. [Checklist: Adding a New Node](#2-checklist-adding-a-new-node)
|
||||
3. [Checklist: Adding a New Setting / Feature Flag](#3-checklist-adding-a-new-setting--feature-flag)
|
||||
4. [Checklist: Updating Dependencies](#4-checklist-updating-dependencies)
|
||||
5. [Cross-Platform & Safe Coding Standards](#5-cross-platform--safe-coding-standards)
|
||||
6. [User Data & Storage Architecture (`ComfyUI/user/default/SaturnNodes/`)](#7-user-data--storage-architecture-comfyuiuserdefaultsaturnnodes)
|
||||
7. [Git, Branching & Commit Conventions](#6-git-branching--commit-conventions)
|
||||
|
||||
---
|
||||
|
||||
## 1. Architecture & Directory Layout
|
||||
|
||||
```
|
||||
ComfyUI-SaturnNodes/
|
||||
├── __init__.py # Node registration (NODE_CLASS_MAPPINGS), settings routes (/saturnnodes/*)
|
||||
├── pyproject.toml # Package metadata, dependencies & ComfyUI Registry configuration
|
||||
├── requirements.txt # Pip dependencies (auto-installed by ComfyUI Manager)
|
||||
├── .env.example # Template of all supported environment variables
|
||||
├── .gitignore # Exclusion rules for caches, states, and credentials
|
||||
├── .comfyignore # Files excluded from the ComfyUI Registry package archive
|
||||
├── nodes/ # Python backend node implementations
|
||||
│ ├── aspect_ratio.py # Text & preview aspect ratio calculation nodes
|
||||
│ ├── assets_restore.py # Startup output-folder history restoration
|
||||
│ ├── auto_watcher.py # Automated folder watch image loaders
|
||||
│ ├── decision_node.py # Inline pause/gate decision node & notifications
|
||||
│ ├── generation_times.py # Prompt duration recording for restored asset history
|
||||
│ ├── image_loader.py # Visual image browser & EXIF/PNG metadata parser
|
||||
│ ├── load_recent.py # Recent outputs loader node
|
||||
│ ├── local_runner.py # Confined local script execution node & authorization routes
|
||||
│ ├── lora_finder.py # Text-based prompt LoRA scanner & patcher
|
||||
│ ├── lora_loader.py # Folder & visual LoRA loaders + Civitai/TMDB scrapers
|
||||
│ ├── preview_latent.py # Live latent WebSocket preview node
|
||||
│ ├── prompt_counter.py # Multiline prompt counter node
|
||||
│ ├── prompt_iterator.py # Multiline prompt queue batch popper
|
||||
│ ├── queue_control.py # Pause Queue, Persistent Queue hooks, and server routes
|
||||
│ ├── stereogram.py # Stereoscopic 3D / cross-eye generator nodes
|
||||
│ ├── text_replacer.py # Multiline regex & list text replacer node
|
||||
│ ├── text_split.py # Forward/backward & regex text splitter node
|
||||
│ ├── tray_icon.py # OS system tray icon manager & status indicator
|
||||
│ └── utils.py # Shared path sanitization, string parsing, and formatting helpers
|
||||
├── web/ # Frontend extensions & static assets
|
||||
│ ├── saturnnodes_colors.js # Node theme colors (LiteGraph canvas & Vue V2)
|
||||
│ ├── saturnnodes_settings.js# Native ComfyUI Settings panel registrations
|
||||
│ ├── pause_queue.js # Top toolbar Pause/Continue button group & dropdown
|
||||
│ ├── pause_queue.css # Toolbar styling
|
||||
│ ├── persistent_queue.js # Queue crash-recovery client ownership claimer
|
||||
│ ├── batch_queue.js # 1D git-graph batch queue visualizer
|
||||
│ ├── floating_latent_dock.js# Floating live latent preview dock
|
||||
│ ├── preview_node.js # Live latent canvas preview widget
|
||||
│ ├── logo.png # Registry icon (400x400)
|
||||
│ └── js/ # Node-specific UI widgets (visual pickers, buttons, aspect preview)
|
||||
│ ├── auth_helper.js # CSRF token helper shared by all frontend extensions
|
||||
│ ├── copy_prompt.js
|
||||
│ ├── decision_node.js
|
||||
│ ├── folder_lora_loader.js
|
||||
│ ├── image_loader.js
|
||||
│ ├── local_runner.js
|
||||
│ ├── power_sidebar.js
|
||||
│ ├── preview_aspect_ratio.js
|
||||
│ ├── preview_manager.js
|
||||
│ ├── prompt_counter.js
|
||||
│ ├── prompt_iterator.js
|
||||
│ └── stereogram.js
|
||||
├── README.md # Main documentation & complete node reference
|
||||
├── WALKTHROUGH.md # End-user guide for workflows and features
|
||||
└── CHANGELOG.md # Version release history
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Checklist: Adding a New Node
|
||||
|
||||
When implementing a new custom node, you **must complete all 5 steps**:
|
||||
|
||||
### Step 2.1: Python Backend (`nodes/<module>.py`)
|
||||
- [ ] Define standard class attributes:
|
||||
- `INPUT_TYPES(cls)`: Return dict with `required`, `optional`, and `hidden` widgets. Always specify defaults, min/max/step for numbers.
|
||||
- `RETURN_TYPES`: Tuple of output types (e.g. `("IMAGE", "STRING")`).
|
||||
- `RETURN_NAMES`: Tuple of human-readable output labels.
|
||||
- `FUNCTION`: Exact method name to execute.
|
||||
- `CATEGORY`: Categorized under `"🪐 SaturnNodes/Loaders"`, `"🪐 SaturnNodes/Utils"`, `"🪐 SaturnNodes/Automation"`, `"🪐 SaturnNodes/Previews"`, or `"🪐 SaturnNodes/3D & Previews"`.
|
||||
- `DESCRIPTION`: Clear multi-line string explaining node functionality and parameters.
|
||||
- [ ] Use `sanitize_folder_path(folder_input, default_dir)` from `nodes.utils` for any path/directory arguments (handles Windows/Linux slashes, wildcards `*`, and relative paths cleanly).
|
||||
- [ ] Ensure dummy tensors are valid 4D float32 batches `torch.zeros((1, 64, 64, 3), dtype=torch.float32)` for empty image states.
|
||||
- [ ] Use `comfy.model_management.throw_exception_if_processing_interrupted()` inside long loops.
|
||||
|
||||
### Step 2.2: Export & Map in `__init__.py`
|
||||
- [ ] Import the node class into [`__init__.py`](./__init__.py).
|
||||
- [ ] Register in `NODE_CLASS_MAPPINGS`:
|
||||
```python
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"YourNewNodeName": YourNewNodeClass,
|
||||
# ...
|
||||
}
|
||||
```
|
||||
- [ ] Register in `NODE_DISPLAY_NAME_MAPPINGS` with the Saturn prefix and emoji:
|
||||
```python
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"YourNewNodeName": "🪐 🏷️ Your Display Name",
|
||||
# ...
|
||||
}
|
||||
```
|
||||
- [ ] Avoid renaming or replacing existing node class names; if unavoidable, migrate known workflow files to the new class name instead of keeping permanent legacy aliases (see §7).
|
||||
|
||||
### Step 2.3: Frontend Extensions & Colors (`web/`)
|
||||
- [ ] If the node needs custom canvas/DOM widgets, implement the extension under `web/js/<feature>.js`.
|
||||
- [ ] In `nodeCreated(node)` or `beforeRegisterNodeDef(nodeType, nodeData)`, **always check both `node.comfyClass` and `node.type`**:
|
||||
```javascript
|
||||
const isTargetNode = node.comfyClass === "YourNodeName" ||
|
||||
node.type === "YourNodeName";
|
||||
```
|
||||
- [ ] Register node theme colors in [`web/saturnnodes_colors.js`](./web/saturnnodes_colors.js):
|
||||
- The suite uses a single unified Saturn amber theme (`DEFAULT_HEADER_COLOR = "#6a3f20"`, `DEFAULT_BG_COLOR = "#472a15"`), user-overridable via the Custom Node Title/Body Color settings.
|
||||
- Add the new node class key to `SATURN_NODE_TYPES` so the custom colors apply to it.
|
||||
|
||||
### Step 2.4: Documentation (`README.md` & `WALKTHROUGH.md`)
|
||||
- [ ] Add a clean heading section in [`README.md`](./README.md) under `## 📦 Individual Nodes Reference`:
|
||||
- Heading: `### 🪐 🏷️ Display Name — \`ClassName\``
|
||||
- `#### Overview`
|
||||
- `#### Inputs & Widgets`
|
||||
- `#### Outputs`
|
||||
- [ ] If the node introduces a new workflow pattern, add a section in [`WALKTHROUGH.md`](./WALKTHROUGH.md).
|
||||
|
||||
### Step 2.5: Changelog
|
||||
- [ ] Add entry under `### Added` in [`CHANGELOG.md`](./CHANGELOG.md).
|
||||
|
||||
---
|
||||
|
||||
## 3. Checklist: Adding a New Setting / Feature Flag
|
||||
|
||||
When introducing a configurable option or feature toggle:
|
||||
|
||||
### Step 3.1: `.env.example`
|
||||
- [ ] Add the variable with default value and explanation in [`.env.example`](./.env.example):
|
||||
```ini
|
||||
# Brief description of what this setting controls
|
||||
YOUR_NEW_SETTING=true
|
||||
```
|
||||
|
||||
### Step 3.2: Backend Settings Synchronization (`__init__.py`)
|
||||
- [ ] In `get_settings()` route: parse and return the setting.
|
||||
- [ ] In `save_settings()` route: read value from incoming JSON, persist to `.env` file via `new_lines`, update `os.environ`, and trigger any dynamic runtime manager updates.
|
||||
|
||||
### Step 3.3: Frontend Settings Registration (`web/saturnnodes_settings.js`)
|
||||
- [ ] Register via `app.ui.settings.addSetting({...})`:
|
||||
```javascript
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.YourSettingName",
|
||||
name: "🪐 SaturnNodes: Your Setting Display Name",
|
||||
type: "boolean" /* or "text", "combo" */,
|
||||
defaultValue: false,
|
||||
tooltip: "Descriptive tooltip explaining user impact.",
|
||||
onChange(value) {
|
||||
api.fetchApi("/saturnnodes/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ your_setting_key: value })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
```
|
||||
|
||||
### Step 3.4: Documentation
|
||||
- [ ] Update `## ⚙️ ComfyUI Settings Menu Reference` in [`README.md`](./README.md).
|
||||
- [ ] Update `## 4. Settings & Configuration` in [`WALKTHROUGH.md`](./WALKTHROUGH.md).
|
||||
|
||||
---
|
||||
|
||||
## 4. Checklist: Updating Dependencies
|
||||
|
||||
If a feature requires a new Python package:
|
||||
- [ ] Add the package to `dependencies` in [`pyproject.toml`](./pyproject.toml) (ComfyUI Registry packaging).
|
||||
- [ ] Keep [`requirements.txt`](./requirements.txt) in sync for ComfyUI Manager auto-install.
|
||||
- [ ] Check imports gracefully handle missing optional libraries (e.g. `try... except ImportError`).
|
||||
- [ ] **Never** add an `install.py` or any runtime `pip`/`subprocess` installation — the suite is ComfyUI Registry compliant and installs nothing dynamically.
|
||||
|
||||
---
|
||||
|
||||
## 5. Cross-Platform & Safe Coding Standards
|
||||
|
||||
1. **Terminal Console Safety (Windows `cp1252` encoding)**:
|
||||
- **Never include raw unicode emojis in backend `print()` calls** (e.g. write `print("[SaturnNodes] Started")` instead of `print("[SaturnNodes] 🪐 ...")`).
|
||||
- Windows consoles running on default codepages will raise fatal `UnicodeEncodeError: 'charmap' codec can't encode character` if raw emoji bytes hit standard output.
|
||||
2. **Path Sanitization**:
|
||||
- Never assume `/` or `\` separators; always use `os.path.join()`, `os.path.normpath()`, or `sanitize_folder_path()`.
|
||||
3. **OS Platform Detection**:
|
||||
- Use `sys.platform` (`'win32'`, `'darwin'`, `'linux'`) rather than POSIX-only calls like `os.uname()`.
|
||||
4. **Non-Blocking & Daemon Threads**:
|
||||
- All background threads (scraping, tray icons, persistent queue sync) must set `daemon=True` so they do not block ComfyUI shutdown.
|
||||
5. **Bytecode Validation**:
|
||||
- Before committing, always run:
|
||||
```bash
|
||||
python -m compileall .
|
||||
```
|
||||
Ensure 0 syntax errors across all modules.
|
||||
|
||||
---
|
||||
|
||||
## 7. User Data & Storage Architecture (`ComfyUI/user/default/SaturnNodes/`)
|
||||
|
||||
All user-specific runtime states, API credentials, usage counters, and temporary caches **MUST** be stored centrally in `ComfyUI/user/default/SaturnNodes/` using the helper `get_saturnnodes_user_dir()` in `nodes/utils.py`:
|
||||
|
||||
```python
|
||||
from .utils import get_saturnnodes_user_dir
|
||||
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
STATE_FILE = os.path.join(USER_DIR, "your_state.json")
|
||||
```
|
||||
|
||||
### Stored User Files:
|
||||
- `ComfyUI/user/default/SaturnNodes/.env` — API keys and persistent settings
|
||||
- `ComfyUI/user/default/SaturnNodes/lora_usage.json` — LoRA selection counts and usage badges
|
||||
- `ComfyUI/user/default/SaturnNodes/image_prompts_cache.json` — Cached positive prompts extracted from images
|
||||
- `ComfyUI/user/default/SaturnNodes/generation_times.json` — Recorded prompt durations and output filenames for restored asset history
|
||||
- `ComfyUI/user/default/SaturnNodes/failed_scrapes.json` — Failed Civitai/TMDB scraping attempts cache
|
||||
- `ComfyUI/user/default/SaturnNodes/prompt_iterator_state.json` — Current iteration pointers and popped prompt queues
|
||||
- `ComfyUI/user/default/SaturnNodes/persistent_queue.json` — Real-time queue crash recovery state
|
||||
|
||||
### Directory Rules & Compatibility Policy:
|
||||
- All state, caches, and credentials live in `ComfyUI/user/default/SaturnNodes/` only. No legacy directories are read, migrated, or referenced.
|
||||
- No pre-rename compatibility shims exist or may be reintroduced (old HTTP routes, old setting keys/IDs, old localStorage keys, old headers, old user directories, old node class-name aliases). Old workflow files are migrated to the current node names instead.
|
||||
|
||||
---
|
||||
|
||||
## 6. Git, Branching, Versioning & Release Strategy (MANDATORY)
|
||||
|
||||
### 6.1 Version-Named Branching Workflow
|
||||
From version `v2.5.0` onward, all development work in this repository strictly adheres to a version-scoped branching model:
|
||||
1. **Active Work Occurs ONLY in the Current Version Branch**:
|
||||
- All active development, bug fixes, UI improvements, and code changes **MUST** occur on a dedicated branch named directly after the target version prefixed with `v` (e.g. `v2.5.0`).
|
||||
- Git and GitHub fully support dots in branch names (`v2.5.0`, `v2.6.0`, `v3.0.0`).
|
||||
- **STRICTLY PROHIBITED**: Never commit or push directly to `main` during active development.
|
||||
2. **Strict User-Approval Gate for Merges & Releases**:
|
||||
- Merging to `main` is **STRICTLY GATED** by explicit user approval.
|
||||
- **NEVER** merge a version branch into `main` unless the user explicitly commands: *"merge to main"* or *"make a release"*.
|
||||
3. **Release & Tagging Execution Steps (Upon User Command)**:
|
||||
When explicit permission to release is granted:
|
||||
- **Merge**: Merge the version branch (e.g. `v2.5.0`) into `main`.
|
||||
- **Tag**: Create an annotated Git tag prefixed with `v` (e.g. `git tag -a v2.5.0 -m "v2.5.0 ..."`).
|
||||
- *Rule*: **ALWAYS** use the `v` prefix (`v2.5.0`, `v2.6.0`). **NEVER** create bare tags without `v` (like `2.5.0`).
|
||||
- **GitHub Release**: Publish or update the GitHub Release attached to that `vX.Y.Z` tag.
|
||||
- *Rule*: Every release **MUST** include a clear, descriptive title (e.g. `v2.5.0 - <Feature/Theme Summary> 🪐`) and a concise markdown bullet-point explanation of what changed and what was fixed (matching the standard of `v2.3.2` and `v2.4.0`), never just bare automated compare links.
|
||||
- **Next Version Transition**:
|
||||
- Immediately checkout a new branch from `main` named after the next target version (e.g. `git checkout -b v2.6.0` or `git checkout -b v3.0.0`).
|
||||
- Update version identifiers in `pyproject.toml`, `__init__.py`, `CHANGELOG.md`, and frontend info to the new development version.
|
||||
- Continue all subsequent work exclusively on the new version branch.
|
||||
|
||||
### 6.2 Excluded Files & Privacy
|
||||
Never commit sensitive keys, personal paths, runtime states, or caches:
|
||||
- `.env`, `.env.local`
|
||||
- `persistent_queue.json`, `lora_usage.json`, `failed_scrapes.json`, `image_prompts_cache.json`
|
||||
- `user/` directory
|
||||
- Personal user names, machine paths (`C:\Users\...`), or private model names.
|
||||
|
||||
### 6.3 Commit Messages
|
||||
Use atomic, descriptive English commit headers:
|
||||
- `Feat: ...` (New features or nodes)
|
||||
- `Fix: ...` (Bug fixes, UI repairs)
|
||||
- `Docs: ...` (Documentation, README, Walkthrough updates)
|
||||
- `Refactor: ...` (Code restructuring or cleanup)
|
||||
- `Release: ...` (Consolidated version release merge)
|
||||
|
||||
### 6.4 Registry Publishing (ComfyUI Registry)
|
||||
- Registry node name: `saturnnodes` (Publisher `KOFiblto`, `DisplayName = ComfyUI-SaturnNodes`). **Published versions are immutable** — always bump the version before publishing a change.
|
||||
- `.github/workflows/publish_action.yml` publishes automatically when `pyproject.toml` changes on `main` (and can also be triggered manually via workflow dispatch). It requires the `REGISTRY_ACCESS_TOKEN` repository secret.
|
||||
- `.comfyignore` excludes development files (tests, CI, contributor docs) from the published package archive.
|
||||
- Icon (`web/logo.png`) must be square, 400x400 or smaller.
|
||||
+302
-2
@@ -1,10 +1,310 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to `ComfyUI-FlowControl` will be documented in this file.
|
||||
All notable changes to `ComfyUI-SaturnNodes` will be documented in this file.
|
||||
|
||||
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
|
||||
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
||||
|
||||
## [2.6.3] - 2026-10-30 [Stable]
|
||||
|
||||
### Added
|
||||
- **Restore Regression Test Matrix**: 64 new deterministic Python tests and 3 new Node fake-DOM tests covering DOM injection, cross-day timestamps, equal-mtime ordering, restore limits, corrupt images, batches, cancellation, retries, corrupt state, and recovery. The existing `test_batch_queue_dom.mjs` is now run by the suite as well.
|
||||
- **Queue Bleed Regression Tests**: `test_batch_queue_snapshot_dom.mjs` models two Queue clicks with an in-flight submission and asserts prompt text/LoRA freezing with live seed merging, and `test_prompt_iterator.py` covers linked-prompt precedence over stale local text.
|
||||
|
||||
### Changed
|
||||
- **Deterministic Restore Ordering**: Output scanning and the Visual Image Loader now tie-break equal modification times by relative path, so bulk-copied or archive-extracted files always restore in the same order and receive stable history numbering.
|
||||
- **Restored Queue Identity Repair**: `persistent_queue.json` entries are normalized on load and before restore. The top-level `prompt_id` is authoritative, `item[1]` is re-synced before queueing, duplicates collapse to the last occurrence, and malformed entries are dropped instead of crashing recovery.
|
||||
- **Synthetic Restored Queue Timestamps**: Restored items receive strictly increasing `create_time` values ending just before "now", preserving their saved queue order instead of sharing one timestamp.
|
||||
- **Best-Effort Asset DB Ingestion**: A failing `app.assets` registration no longer aborts a successful history restore.
|
||||
- **Serialized Restores**: Startup restore and manual "Restore Now" requests are serialized by a restore lock so overlapping calls cannot interleave history writes or asset ingestion.
|
||||
|
||||
### Fixed
|
||||
- **Restored Duration Validation**: `NaN`, infinite, negative, and non-numeric durations are rejected when recorded and when read, so they can no longer poison restored timestamps or produce invalid JSON.
|
||||
- **Generation Times Persistence**: `generation_times.json` saves now sanitize corrupt entries, use unique temp files, retry transient Windows `os.replace` locks, serialize with `allow_nan=False`, and clean up leftover temp files.
|
||||
- **Cancel/Wipe Recovery Loss**: Persistent queue entries are removed only after ComfyUI confirms the cancel or wipe; a failed operation keeps the crash-recovery copy intact.
|
||||
- **Restore Robustness**: Missing queue mutexes, non-dict histories, history-read failures, malformed retry counters, and invalid restore limits are handled without aborting recovery.
|
||||
- **EXIF / A1111 Metadata**: EXIF UserComments with BOM/NUL padding and byte-valued A1111 `parameters` now decode to clean prompt text.
|
||||
- **Batch Color Injection**: Batch colors from `localStorage`/the server are restricted to plain hex values before entering SVG markup; invalid colors fall back to the palette.
|
||||
- **Batch Counter & Metadata Corruption**: Negative persisted batch counters clamp into the palette range, metadata pruning uses `queuedAt` recency, unknown segment types normalize safely, out-of-range item indexes clamp to the batch bounds, and older server snapshots can no longer overwrite newer local batch metadata.
|
||||
- **Concurrent Batch Item Indexes**: Concurrent submissions reserve unique `itemIndex` values through a captured batch context and `inFlight` counter, so two prompts can no longer claim the same index or clear the batch early.
|
||||
- **Recent Output Loading**: `Load Recent Outputs` clamps invalid `amount`/`index` values, tie-breaks equal mtimes by path, and closes image handles after loading.
|
||||
- **Registry Publisher ID**: `PublisherId` now matches the registry's case-sensitive publisher id `kofiblto` (the previous `KOFiblto` resolved to 404 on `api.comfy.org`), and the publish workflow fails fast with a readable message when `REGISTRY_ACCESS_TOKEN` is missing instead of running a full install first.
|
||||
- **Sequential LoRA Cycle Reordering**: The Sequential cycle now persists the selection order it started with instead of indexing the live (reorderable) list, so deselecting and reselecting LoRAs can no longer repeat one LoRA and skip another from the same state bucket.
|
||||
- **Missing LoRA Selection Mode Default**: `load_lora` and `IS_CHANGED` now fall back to `Sequential` (matching `INPUT_TYPES` and the UI) instead of `All`, so API prompts without `_selection_mode` cycle one LoRA per prompt instead of stacking every selected LoRA into one render.
|
||||
- **Click-to-Queue Prompt Bleed**: Queue clicks made while another batch is still being submitted now serialize the graph at click time. Text, LoRA selections and other node inputs edited afterwards can no longer bleed into the queued prompt, and single-item queues get the same protection as multi-item batches while random seeds still follow the live prompt.
|
||||
- **Prompt Iterator Linked Input**: A connected `prompt` input now takes precedence over the node's local multiline `text`, so stale local text can no longer be replayed instead of the upstream *Prompt Counter* output.
|
||||
- **LoRA Picker Console Flood**: The Visual LoRA Loader now reads its thumbnail setting once per picker render and without the deprecated `getSettingValue(id, default)` argument, so opening a large LoRA folder no longer emits a console warning for every tile.
|
||||
|
||||
---
|
||||
|
||||
## [2.6.2] - 2026-10-24 [Stable]
|
||||
|
||||
### Added
|
||||
- **Per-Resolution Dock Memory & Drag Resize Handles**: The Floating Latent Dock and the queue overlay now remember their size per resolution/monitor and expose on-screen drag resize handles.
|
||||
- **Visual LoRA Loader Maintenance Controls**: New `↺ Counter`, `🧹 Usage` and `🗑 Cache` reset buttons on the Visual LoRA Loader node (each with a confirmation prompt) to reset the per-node Sequential/Random LoRA cycle and prompt→LoRA assignments, clear usage badges, or delete cached thumbnails.
|
||||
- **ComfyUI Registry Packaging**: Complete registry metadata (`pyproject.toml` with `requires-comfyui`, classifiers and optional extras), a Saturn icon (`web/logo.png`), a `.comfyignore` manifest, and an automated `publish_action.yml` workflow.
|
||||
- **Curated GitHub Releases**: `release.yml` now publishes the matching `CHANGELOG.md` section as the GitHub Release body instead of raw auto-generated compare links.
|
||||
- **Real Generation Times on Restored Assets**: SaturnNodes now records prompt execution duration and output filenames in `generation_times.json`; restored history entries expose real start/success timestamps and the stock duration display when available. Older images keep their file mtime for sorting but show no fabricated `0.00s` duration.
|
||||
- **Prompt Iterator Tracking Modes & Exhaustion Output**: New `track_by` choice (`Text & Node` / `Node Only`) and `exhausted` boolean output for `Sequential (Stop on End)`.
|
||||
- **Persistent Pause State**: Optional `PERSIST_PAUSE_STATE` setting restores the last manual Pause/Continue state and pause mode after a restart.
|
||||
- **Queue Restore Retry Limit**: Optional `PERSISTENT_QUEUE_MAX_RETRIES` setting controls how many later launches may retry failed queue items (default `0`).
|
||||
- **Assets Restore Controls**: New recursive-folder scan toggle, generation-time display toggle, and a **Restore Now** settings button.
|
||||
- **Visual Image Loader Metadata Limit**: New `IMAGE_LOADER_METADATA_LIMIT` setting (default `128`) controls how many newest images are scanned for prompts.
|
||||
|
||||
### Changed
|
||||
- **Free-Form Dock Resize**: Dragging the Floating Latent Dock's corner now resizes width and height independently (no longer locked to the latent's aspect ratio) within sane 100–2000px bounds. The drag starts from the dock's current on-screen size (no jump to the larger saved settings box), shows the free-form box live, writes the dragged dimensions back into the `Floating Dock Max Width/Height` settings (now 1px-stepped sliders), and snaps the dock to the exact latent fit on release.
|
||||
- **Route Namespacing**: All HTTP routes live exclusively under `/saturnnodes/*`; the transitional pre-2.6 route aliases (`/pause_queue/*`, `/persistent_queue/*`, `/image_loader/*`, `/folder_lora_loader/*`) were removed entirely.
|
||||
- **Compatibility Surface Reduction**: All remaining pre-rename compatibility layers were purged — duplicate legacy route handlers, frontend settings migration and legacy settings-ID lookups, legacy localStorage/storage keys, the legacy CSRF header fallback, the legacy user-directory auto-migration, and all superseded node class-name aliases. No backward-compatibility surface remains.
|
||||
- **Persistent Queue Schema Versioning**: `persistent_queue.json` now uses a top-level `{"schema_version": 1, "items": [...]}` envelope. Legacy bare-list files are migrated atomically on first load, files written by a newer SaturnNodes are kept read-only when downgrading, and transient Windows sharing violations during the atomic replace are retried.
|
||||
- **Naming Consistency**: Unified `ComfyUI.SaturnNodes.*` frontend extension names, `saturnnodes-*` CSS prefixes, and `[SaturnNodes ...]` console log tags.
|
||||
- **Documentation Rebrand**: `WALKTHROUGH.md`, the README settings reference, `.env.example`, `AGENTS.md` and `SECURITY.md` were fully rebranded and reconciled with the real setting labels, widget names and environment variables.
|
||||
- **Registry Identity**: Package name set to `saturnnodes` (per registry naming best practices); version bumped to `2.6.2`.
|
||||
|
||||
### Fixed
|
||||
- **Publishing Compliance**: Removed `install.py` (no runtime `pip` execution, matching the documented "Zero Runtime Package Installs" claim) and corrected a stale README claim about legacy route aliases.
|
||||
- **Test Isolation**: The test suite now redirects every module-level SaturnNodes state path (persistent queue, LoRA usage/cycle state, caches, iterator state, generation times) to a disposable temporary user directory. Running tests while ComfyUI is active can no longer overwrite live user data.
|
||||
- **Audit Remediation Hardening**: Full audit remediation pass - pause queue stall/deadlock fixes, restart verification hardening, atomic `.env` writes with proper error responses, debug endpoint guards, subgraph-safe lookups, runner pipe streaming, IPv4-mapped loopback handling and more (160 unit tests passing).
|
||||
- **DOM Reuse & Recycled Widget Hardening**: Fixed all eight stale-DOM audit findings. Power-popup document listeners now die with their popup and a detached popup can no longer swallow a reopen. Copy Prompt revalidates the interacted card by asset identity, rejects detached/recycled cards, re-resolves the image at click time, and rebinds legacy hover-overlay buttons. The floating latent dock observes deep minimap/toolbar/sidebar layout changes. Remounted node widgets (live preview image, prompt counter badge, LoRA count badge) re-resolve their live DOM element instead of writing to detached nodes. Batch lines are painted when a legacy v1 queue table mounts with its rows already built, duplicate `data-job-id` rows no longer skew batch grouping, and the saved queue height is re-applied to `.comfy-list` elements mounted after setup.
|
||||
- **Restored Queue Batch-Line Repaint**: Queue cards reused by ComfyUI's virtualized Vue list could keep the previous batch's colored line when the next batch started, making the running item appear to jump from the pending color to the previous batch color. Batch lines now repaint on `data-job-id` attribute changes and during a short post-event refresh window; unchanged SVG lines are skipped to avoid unnecessary DOM churn.
|
||||
- **Floating Latent Dock Resize Crash**: Fixed a `this.renderImage is not a function` error thrown while dragging the dock's resize handles or when the screen profile changed; the resize paths now invoke the dock's actual render routine.
|
||||
- **Queue Overlay Resize Lag**: Resizing the queue overlay no longer rewrites a stylesheet on every pointermove (which forced a full-document style recalculation, layout and paint churn on every frame); heights are now applied as scoped, rAF-throttled inline styles.
|
||||
- **Dock Drag Performance & Stuck Drags**: Dragging the dock no longer rebuilds the full-resolution canvas backing store (GPU texture/layer churn, ~74% GPU load during drags) on every pointermove — only CSS sizes change while dragging, with one full-quality redraw on release. Pointer capture plus `pointercancel`/window-blur guards fix drags that "never stopped" when the pointer was released outside the window.
|
||||
- **Prompt Iterator Reset Clarity**: The Prompt Queue Iterator's `Reset Counter (0)` button now confirms before resetting and explains that it resets the run pointer to the first prompt without wiping the queue state file (use Settings → `Reset All Queues` for a full wipe).
|
||||
- **Prompt Iterator Correctness & Concurrency**: Hard resets now clear crash-recovery assignments and shuffle permutations; state keys include splitter/mode configuration; `Random (Cycle)` remaining text follows the real shuffle order; `Random (Keep)` no longer reports a fake remaining ordered list; state reads use an mtime cache behind a lock; global state pruning keeps the most recently used entries.
|
||||
- **LoRA Loader State Growth & Performance**: Prompt→LoRA assignment maps and global cycle state are now bounded by recency; selection mapping is resolved once per run; usage counters are batched in memory and flushed atomically; No-Replace cycles avoid an immediate repeat while still visiting every LoRA; Civitai SHA256 hashes are cached by mtime/size; `TextLoraFinder` scans the prompt once instead of once per installed LoRA.
|
||||
- **Persistent Queue I/O Storm**: `persistent_queue.json` writes are debounced, compact and fsynced instead of rewriting the full file for every queued/removed item; batch metadata and recently-finished prompt IDs are capped to prevent unbounded growth; late batch syncs can no longer resurrect finished prompts.
|
||||
- **Assets Restore Blocking & Coverage**: Output scanning, image metadata parsing and asset DB registration now happen outside `prompt_queue.mutex`; restore uses a bounded top-N scan and a single DB session; JPEG/WebP EXIF metadata and A1111/Forge `parameters` prompts are supported; debug reports include elapsed time, scanned count, metadata hits and generation-time matches.
|
||||
- **Visual Image Loader Responsiveness**: Folder scans and thumbnail generation run off the asyncio event loop; thumbnails are cached on disk by path/mtime/size.
|
||||
- **Settings/.env Divergence**: `/saturnnodes/settings` now returns the complete effective settings surface, the UI seeds from backend values before registering defaults, and parsed `.env` values are cached with mtime invalidation and explicitly invalidated after saves.
|
||||
|
||||
### Removed
|
||||
- `install.py` — dependencies are declared in `pyproject.toml` / `requirements.txt` only, and ComfyUI Manager handles installation.
|
||||
|
||||
---
|
||||
|
||||
## [2.5.0] - 2026-10-10 [Stable]
|
||||
|
||||
### Added
|
||||
- **Floating Latent Live Preview Dock & Core Lightbox Live Inspection**:
|
||||
- Floating live latent preview dock with dynamic aspect ratio preservation and contain-scaling, eliminating letterbox bars and layout jitter.
|
||||
- Minimap tracking and alignment alongside canvas navigation controls.
|
||||
- **Native Core Lightbox Live Inspecting View**: Clicking the floating latent dock opens the live generation directly in ComfyUI's native inspecting lightbox view, updating live in real time as each diffusion sampling step arrives.
|
||||
- Minimalist design: removed redundant "Sampling..." badge overlays for distraction-free monitoring.
|
||||
- **High-Performance LoRA Loader Engine & Viewport WebP Lazy-Loading**:
|
||||
- **Disk-Cached WebP Generation**: High-resolution preview images are downscaled via bicubic fast-WebP encoding and cached in `ComfyUI/user/default/SaturnNodes/cache/lora_thumbs/` (~35KB per card).
|
||||
- **Non-Blocking Background Worker**: Thumbnail processing runs asynchronously on worker threads with `threading.Event`, eliminating canvas UI thread lockups during heavy directory indexing.
|
||||
- **Viewport IntersectionObserver**: Implemented native `IntersectionObserver` on visual LoRA cards so thumbnails only load as they scroll into view.
|
||||
- Added an explicit settings toggle to bypass thumbnail generation if desired.
|
||||
- **Stereoscopic 3D / Cross-Eye Generator (`SaturnStereo3D`)**:
|
||||
- Transform 2D images into stereoscopic 3D pairs for cross-eye viewing, parallel 3D viewers, VR displays, and red-cyan anaglyphs.
|
||||
- Hybrid depth execution: accepts any input `depth_map` (zero downloads, 1ms execution) or automatically falls back to built-in Depth Anything V2 Small ONNX model cached locally, with mathematical gradient fallback when ONNX is absent.
|
||||
- Stereonix-inspired controls: `depth_intensity`, `depth_gamma`, `invert_depth`, `convergence_plane`, `image_zoom`, `pan_horizontal`, `pan_vertical`.
|
||||
- Framing and focus controls: adjustable `gap_spacing`, `outer_border` padding, and high-contrast `alignment_dots`.
|
||||
- **Stereoscopic 3D Live Viewer (`SaturnStereo3DLive`)**:
|
||||
- Interactive live viewer node featuring a real-time **1s countdown timer** on parameter adjustment.
|
||||
- **Mid-Render Interruption**: Modifying sliders during generation immediately sends `/interrupt` to abort obsolete rendering and restarts the 1s countdown.
|
||||
- **In-Memory Depth Map Caching**: Cached depth predictions allow near-instant (~5ms) re-rendering when tweaking intensity, gamma, zoom, pan, and spacing on the same input image.
|
||||
- **Queue Execution Controls & Mid-Generation Pause**:
|
||||
- Added execution interception hook allowing the queue to pause before popping the next item from the queue.
|
||||
- Restored queue prompt `create_time` timestamps are refreshed upon restart to prevent deprioritization in history.
|
||||
- **Modernized Continuous Sliders Settings UI**:
|
||||
- Converted `Visible Queue Items (Height)` and `Live Latent Dock Size` from rigid dropdowns into smooth continuous sliders.
|
||||
- Cleaned settings categories and isolated context menus to prevent bleeding into top PrimeVue application menubars.
|
||||
- **CLI Terminal Branding & Clean Telemetry**:
|
||||
- Clean unified startup banner styled in Saturn amber theme.
|
||||
- Cleaned noisy internal debug logging, displaying concise status lines for node count, persistence status, and storage paths.
|
||||
|
||||
### Fixed
|
||||
- **Persistent Queue Restored Creation Timestamp**:
|
||||
- Fixed an issue where unfinished prompts resumed after a server restart sorted to the bottom of the history and assets list by updating `create_time` to the current timestamp on recovery.
|
||||
- **Live Viewer Widget Value Corruption**:
|
||||
- Fixed an issue where sliding widgets on `SaturnStereo3DLive` set them to `NaN` or `0` due to arrow-function callback scoping, preventing "wrong type" and "missing connection" validation errors.
|
||||
- **Top Application Menubar Injection**:
|
||||
- Fixed an issue where "Copy prompt" appeared in the main application menu (between File and Edit) on PrimeVue Frontend V2 by strictly excluding `.comfy-command-menu`, `.comfy-menu`, `.p-menubar`, and top-level menuitems.
|
||||
- **CI & Environment Resilience**:
|
||||
- Replaced OpenCV dependency with PIL and ensured graceful fallback when `onnxruntime` is not installed.
|
||||
- **Security & Path Confinement**:
|
||||
- Hardened directory traversal guards against UNC paths and relative escapes, passing 119 unit tests.
|
||||
|
||||
---
|
||||
|
||||
## [2.4.0] - 2026-09-24 [Stable]
|
||||
|
||||
### Fixed
|
||||
- **Frontend V1 Sizing & Layout Stability**:
|
||||
- Eliminated rapid runaway downward expansion loop in `VisualLoraLoader` and `ImageLoader` on LiteGraph classic (Frontend V1).
|
||||
- Cleaned duplicate function stubs and unified `computeSize()` and `computeLayoutSize()`.
|
||||
- **Core Stability & Logic**:
|
||||
- Implemented true cycle permutation for `Random (Cycle)` in `PromptQueueIterator`.
|
||||
- Upgraded comma parsing in `MultiTextReplacer` with `csv.reader` to preserve quoted phrases.
|
||||
- Removed legacy pre-rebrand node class mappings.
|
||||
- Added robust exception handling for unterminated regex in text splitting.
|
||||
- Optimized 60fps canvas draw loops by decoupling badge geometry enforcement from `onDrawForeground`.
|
||||
|
||||
---
|
||||
|
||||
## [2.3.2] - 2026-09-22
|
||||
|
||||
### Changed & Rebranded
|
||||
- **Full Suite Rebrand to SaturnNodes (`🪐`)**:
|
||||
- Renamed the entire suite to **`ComfyUI-SaturnNodes`** with signature emoji `🪐`.
|
||||
- Node display names updated with the `🪐` prefix across all 16 custom nodes.
|
||||
- Node category updated to `🪐 SaturnNodes/*` (`Loaders`, `Utils`, `Previews`, `Queue`, `Automation`).
|
||||
- Full backward compatibility: all legacy pre-rebrand class names preserved and aliased alongside the new class names in `NODE_CLASS_MAPPINGS`.
|
||||
- Dual HTTP route registration for full compatibility with existing scripts, extensions, and bookmarks.
|
||||
- Dual WebSocket event emission for full frontend compatibility.
|
||||
- Seamless settings migration: frontend settings and local storage automatically migrated to `SaturnNodes.*` keys with instant fallback.
|
||||
- Automatic user directory migration: `.env` and `prompt_iterator_state.json` migrated smoothly into `user/default/SaturnNodes/`.
|
||||
- Removed third-party PromptBookmarks coupling: cleaned UI bookmarks button from copy prompt node for modularity.
|
||||
|
||||
---
|
||||
|
||||
## [2.3.1] - 2026-09-22
|
||||
|
||||
### Security & Hardening
|
||||
- **Universal CSRF Protection Across All Mutation Routes**:
|
||||
- Protected the prompt iterator reset, clear, and open-file endpoints with `is_authenticated_local_request`.
|
||||
- Protected the decision node endpoint with `is_authenticated_local_request`.
|
||||
- Protected all queue-mutating endpoints (pause controls, persistent queue claim, batch queue sync, assets restore) with `is_authenticated_local_request`.
|
||||
- Updated all frontend callers in `web/` to send the session CSRF header via `authenticatedFetch`.
|
||||
- **SSRF Protection & Scheme Filtering for Preview Scrapers**:
|
||||
- Added `is_safe_external_image_url` validating that downloaded image URLs use standard `https` or `http` schemes.
|
||||
- Strictly blocked internal loopback (`127.0.0.1`, `localhost`), private networks (`10.*`, `172.16.*`, `192.168.*`), and link-local cloud metadata endpoints (`169.254.*`).
|
||||
- **Safe Write & Permission Error Handling**:
|
||||
- Wrapped downloaded preview image disk writes in `try / except (PermissionError, OSError)` with clear logging to prevent crashes on read-only directories or shared network mounts.
|
||||
- **Canonical Confinement for OS Application Launches**:
|
||||
- Verified that `open_prompt_iterator_file` can only open files strictly confined within the authorized ComfyUI user directory (`USER_DIR`) via `os.path.commonpath`.
|
||||
- **Conditional Core Queue Monkey-Patching**:
|
||||
- Guarded `PersistentQueueManager.patch_server()` so core ComfyUI `prompt_queue` methods are not touched when `ENABLE_PERSISTENT_QUEUE=false`.
|
||||
- **Explicit Scripts Confinement Documentation & UI Canvas Indicators**:
|
||||
- Added explicit italic indicator `📂 Script Location: ComfyUI/scripts/` to `RunLocalFileNode` canvas foreground.
|
||||
- Added widget tooltips and settings descriptions clarifying that scripts and working directories must reside exclusively in `ComfyUI/scripts/`.
|
||||
|
||||
---
|
||||
|
||||
## [2.3.0] - 2026-09-20
|
||||
|
||||
### Added
|
||||
- **🍃 ⚡ "Run Local File" Node (`RunLocalFileNode`)**:
|
||||
- Executes local scripts and binaries (`.bat`, `.cmd`, `.ps1`, `.exe` on Windows; `.sh` and executables on Linux/macOS) directly within ComfyUI workflows.
|
||||
- Parameter arguments supporting multi-line or standard quotes (`shlex` tokenized, safe non-shell execution).
|
||||
- Configurable working directory (`cwd`), defaulting to the script's directory if unspecified.
|
||||
- **Synchronous Mode**: Captures `stdout`, `stderr`, `exit_code`, and `success` boolean outputs, with real-time ComfyUI cancel button interrupt handling and configurable execution timeout.
|
||||
- **Asynchronous Mode**: Detaches process to run in the background without blocking the ComfyUI generation queue.
|
||||
- Optional `trigger` wildcard input and `passthrough` wildcard output for sequencing file execution in any pipeline.
|
||||
- **Multi-Layered Security Sandbox**:
|
||||
- Confined strictly to `ComfyUI/scripts/` via canonical `realpath` and `commonpath` validation.
|
||||
- Strict rejection of absolute paths (`C:\`, `/etc/`, `\\server`) and directory traversal sequences (`..`).
|
||||
- Never invokes `shell=True` (eliminating command chaining and command injection vulnerabilities).
|
||||
- Recursive process tree termination on timeout or cancellation (uses `taskkill /F /T` on Windows to prevent orphaned background processes).
|
||||
- **In-Memory Operator Authorization Gate**: Removed `security_consent` from `INPUT_TYPES` so consent never serializes into workflow JSON or dropped image metadata. Requires active canvas button click `[⚡ Authorize Run]` with 5-minute single-use memory token.
|
||||
- Visual node status indicator (`[ARMED]` in emerald green vs `[DISARMED]` in red).
|
||||
- Global `Allow Local File Execution` toggle setting (default `false`).
|
||||
- **🛡️ CSRF Defense & Request Authentication**:
|
||||
- Added `is_authenticated_local_request` validating loopback origin, `Sec-Fetch-Site` header (blocking `cross-site` calls from visited web pages), and `Origin`/`Referer` headers.
|
||||
- Cryptographically secure per-session CSRF token generated via `secrets.token_hex(32)` and validated with `secrets.compare_digest`.
|
||||
- Added frontend auth helper (`web/js/auth_helper.js`) automatically fetching and attaching CSRF tokens to administrative fetch requests.
|
||||
- **⚡ Two-Step Power Action Handshake**:
|
||||
- Server restart and shutdown endpoints require `ALLOW_PROCESS_MANAGEMENT=true` plus a two-step confirmation handshake (requesting a 30-second single-use ticket, then confirming the action to consume and execute it).
|
||||
|
||||
### Fixed & Compliance
|
||||
- **Complete Removal of `ALLOW_ANY_SCRIPT_PATH`**: Eliminated all path escape settings to guarantee execution confinement to `ComfyUI/scripts/`.
|
||||
- **Dynamic Node Count**: Replaced hardcoded node count in startup banner with dynamic `len(NODE_CLASS_MAPPINGS)` (now reporting all 16 registered nodes).
|
||||
- **Universal Synchronized Versioning**: Version `2.3.0` synchronized across `__init__.py`, `pyproject.toml`, the debug export endpoint, `README.md`, and `CHANGELOG.md`.
|
||||
- **Process Management Security Toggle**: Server restart and shutdown controls are disabled by default (`ALLOW_PROCESS_MANAGEMENT=false`). If disabled, attempts to trigger restart/shutdown via the UI or API return HTTP 403 with clear instructions.
|
||||
- **Optional `pystray` Dependency**: Made `pystray` an optional extra (`[project.optional-dependencies] tray = ["pystray"]`) and removed all runtime `pip install` subprocess calls, ensuring zero crashes on headless cloud servers, Docker, or Google Colab.
|
||||
- **Universal Loopback Protection**: Added `is_local_request(request)` validation to all legacy queue endpoints (`/pause_queue/*` and `/persistent_queue/claim`).
|
||||
- **Windows Console Unicode Safety**: Added `UnicodeEncodeError` guard around startup emoji printing for legacy Windows terminals.
|
||||
- **Clean Dynamic Imports**: Replaced static imports of desktop private `app.*` in `assets_restore.py` with dynamic `importlib.import_module` calls.
|
||||
- **Cleaned Dead Code**: Removed orphaned canvas classes from `queue_control.py`.
|
||||
|
||||
### Testing
|
||||
- Comprehensive automated test coverage expanded to **105 unit tests** covering all nodes, CSRF authentication, power action tickets, in-memory runner authorization, and path confinement.
|
||||
|
||||
---
|
||||
|
||||
## [2.2.0] - 2026-09-08
|
||||
|
||||
> ⚠️ **Important Frontend Compatibility Notice**:
|
||||
> The previous release (**v2.1.0**) does **not** support the new ComfyUI frontend package (`@comfyorg/comfyui-frontend` / modern Vue virtual grid).
|
||||
> **v2.2.0 introduces complete support for the new frontend package**, restoring action buttons (Copy Prompt, Save to Prompt Bookmarks, Inspect Asset zoom) directly on modern virtual grid cards, alongside the 1D Git-Graph batch queue visualizer and per-node zoom settings!
|
||||
|
||||
### Added
|
||||
- **🌐 Full Modern ComfyUI Frontend Package Support**:
|
||||
- Full compatibility with `@comfyorg/comfyui-frontend` modern virtual grid and card DOM structures.
|
||||
- Seamless overlay action bar button groups for modern asset cards with dynamic rounded pill styling (`rounded-l-lg`, `rounded-r-lg`, `rounded-none`, `border-r`).
|
||||
- **🔍 Toggleable "Inspect Asset" Button (`04_EnableInspectAssetButton`)**:
|
||||
- Restores the direct "Inspect asset" (zoom in) button directly onto image cards in the Assets pane hover action bar, placed seamlessly alongside Download, Copy Prompt, and Bookmarks (restoring access since newer ComfyUI versions tucked it away behind the 3-dots submenu). Default is `false` (opt-in).
|
||||
- **🔖 Robust Inline SVG Icons for "Save to Prompt Bookmarks"**:
|
||||
- Replaced runtime Tailwind icon class references with crisp inline SVG ribbons, fixing an issue in modern `comfyui_frontend_package` where the bookmark icon rendered with 0 dimensions (empty space).
|
||||
- **🎨 1D Git-Graph Batch Queue Visualizer**:
|
||||
- Automatically identifies workflows queued together as batches in both ComfyUI Frontend V2 (`JobAssetsList`) and Classic V1 queue list.
|
||||
- Purely graphical 1D colored line along the left edge of queue items without any text or number clutter.
|
||||
- **36 Curated Cycling Colors**: Sequential assignment ensures adjacent batches have distinct, high-contrast colors.
|
||||
- **Start / Middle / End / Single Segments**:
|
||||
- First item in batch segment curves inward at the top (`┌`) and continues straight down.
|
||||
- Middle items connect with a continuous straight line (`│`).
|
||||
- Last item in batch segment comes straight from the top and curves inward at the bottom (`└`).
|
||||
- Single-item batches curve inward at both top and bottom (`(`).
|
||||
- **In-between Infiltration Handling**: When another prompt or batch is queued in-between items of an existing batch, the interrupted batch retains its open straight ends (no false start/end curves), while the inserted item is cleanly enclosed in its own bracket.
|
||||
- **PersistentQueue Interoperability**: Operates client-side via `localStorage` when standalone, and synchronizes with `PersistentQueueManager` so batch relationships restore automatically upon server restart.
|
||||
- **Settings Toggle**: Added an enable/disable batch queue setting in the ComfyUI Settings menu.
|
||||
|
||||
### Fixed
|
||||
- **📐 Independent Per-Node `tile_size` in Visual Loaders**:
|
||||
- Decoupled `tile_size` zoom levels from global browser `localStorage` in `VisualLoraLoader` (V1 & V2) and `VisualImageLoader`.
|
||||
- Tile zoom values are now persisted directly in `node.properties["tile_size"]` within the workflow JSON, allowing multiple loader nodes on the canvas to maintain independent, custom tile sizes across browser reloads.
|
||||
- **🛡️ Queue Container Conflict Prevention**:
|
||||
- Hardened image card button injection guards to avoid injecting into queue item elements or small icon previews in modern virtual grids.
|
||||
|
||||
### Security & Registry Compliance
|
||||
- **Local-Only Route Enforcement**: All server control endpoints (power actions, settings updates, queue sync) now strictly verify loopback origin (`127.0.0.1` / `::1`), rejecting remote calls with `403 Forbidden`.
|
||||
- **Directory Traversal Protection**: Image loading (`VisualImageLoader`) and thumbnail endpoints are strictly confined to ComfyUI `input`, `output`, and `temp` directories.
|
||||
- **Packaging & Installation Standards**: Removed runtime `install.py` in favor of standard `requirements.txt` (`Pillow`, `piexif`, `numpy`, `pystray`), and removed `"torch"` from `pyproject.toml` to prevent host CUDA environment corruption.
|
||||
- **Opt-In Preview Scraping**: Civitai and TMDB network scraping defaulted to disabled (`ENABLE_CIVITAI_SCRAPING=false`, `ENABLE_TMDB_SCRAPING=false`).
|
||||
|
||||
### Changed
|
||||
- **Text Aspect Ratio Finder Whitelist / Open Mode**:
|
||||
- When `aspect_ratios` is populated (e.g. `16:9, 9:16, 1:1`), strictly matches only those configured ratios in prompt text, falling back to default if prompt contains non-whitelisted ratios.
|
||||
- When `aspect_ratios` is left empty (`""`), enters open mode and accepts any valid aspect ratio found in text.
|
||||
|
||||
---
|
||||
|
||||
## [2.1.0] - 2026-08-24
|
||||
|
||||
### Added
|
||||
- **Centralized ComfyUI User Directory (`ComfyUI/user/default/SaturnNodes/`)**:
|
||||
- All runtime JSON state (`lora_usage.json`, `lora_loader_state.json`, `image_prompts_cache.json`, `failed_scrapes.json`, `prompt_iterator_state.json`, `persistent_queue.json`, `.env`) now strictly lives inside the ComfyUI user folder rather than polluting the node directory.
|
||||
- **Structured ComfyUI V2 Settings Menu**:
|
||||
- Restructured settings into 6 explicit subcategories with clear titles, info tooltips, and real styled action buttons (`🗑️ Reset Scrapes Cache`, `🔄 Reset All Queues`).
|
||||
- **Prompt Queue Iterator Live Sync**:
|
||||
- Added interactive Reset widget button directly onto the canvas node and WebSocket live synchronization for multiline prompt text as items are popped.
|
||||
- **Cross-Platform Matrix CI**:
|
||||
- Automated GitHub Actions unit testing across Python 3.10, 3.11, 3.12 on both Ubuntu and Windows with 56 automated unit tests.
|
||||
|
||||
### Fixed
|
||||
- Fixed empty folder matching edge case in `BackToPlaceholder` (`undo_placeholder.py`).
|
||||
- Fixed XML entity escaping for native Windows Toast Notifications in the decision node.
|
||||
- Removed accidental scrape cache reset during minor settings updates in `__init__.py`.
|
||||
- Removed obsolete `playwright` dependency from `pyproject.toml` and `install.py`.
|
||||
|
||||
---
|
||||
|
||||
## [2.0.0] - 2026-08-20
|
||||
|
||||
### Added
|
||||
- **🍃 ⭐ Favorite Prompts**: Preview and loader system for saved favorite prompts and generation presets with subfolder categorizing.
|
||||
- **🍃 📐 Text Aspect Ratio Finder & Preview**: Dynamic aspect ratio parser (e.g. `16:9`, `2:3`) with resolution calculation and visual aspect box preview.
|
||||
- **🍃 🔎 Text LoRA Finder & Loader**: Automatic text prompt scanning for LoRA names with model/clip loading and pretty name formatting.
|
||||
- **🍃 🔄 Prompt Queue Iterator**: Multiline batch prompt queue iterator with Pop, Cycle, and Random modes.
|
||||
- **🍃 🔤 Multi Text Replacer**: High-performance multi-target find & replace node with protection against recursive replacement loops.
|
||||
- **🍃 ✂️ Text Split**: Split strings into two outputs using literal delimiters or regex patterns (forward/backward).
|
||||
- **🍃 ⏸️ Decision**: Interactive mid-workflow pause node with UI buttons (Continue / Cancel / Stop) and desktop notifications.
|
||||
- **🍃 🔔 System Tray Integration**: OS notification area tray icon with real-time queue status colors and outside-browser queue controls.
|
||||
- **🍃 🖼️ Assets & History Restore**: Automatic restoration of recent generations into the ComfyUI Assets / History pane on startup.
|
||||
|
||||
---
|
||||
|
||||
## [1.0.0] - 2026-07-24
|
||||
|
||||
### Added
|
||||
@@ -13,7 +313,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
- **Visual LoRA Loader**: Folder-filtered LoRA selector with formatted pretty names, usage tracking, and Civitai SHA256 API preview thumbnail downloading.
|
||||
- **Visual Image Loader**: Folder image picker with PNG parameters & EXIF user comments metadata parsing.
|
||||
- **Auto Watcher**: Folder watcher with non-blocking check (`wait_for_image = False`) and blocking poll (`wait_for_image = True`).
|
||||
- **Undo Placeholder**: Prompt placeholder restoration tool for LoRA/celebrity names.
|
||||
- **Undo Placeholder**: Prompt placeholder restoration tool for LoRA/subject names.
|
||||
- **Recent Outputs**: Output image loader by recent step index for assets/sidebar integration.
|
||||
- **Live Latent Preview**: Real-time sampler latent preview canvas node renderer.
|
||||
- **Settings Panel**: Native ComfyUI Settings menu controls for API keys and UI toggles.
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# AI Agent Guidelines
|
||||
|
||||
All AI agents (Claude, Gemini, Antigravity, Cursor, Copilot) contributing to this repository must follow the instructions in:
|
||||
- [`AGENTS.md`](./AGENTS.md) — Standard Operating Procedure & Architecture
|
||||
- [`CONTRIBUTING.md`](./CONTRIBUTING.md) — Git Workflow & Contribution Rules
|
||||
+79
-14
@@ -1,19 +1,84 @@
|
||||
# Contributing to ComfyUI-FlowControl
|
||||
# 🤝 Contributing to ComfyUI-SaturnNodes
|
||||
|
||||
Thank you for your interest in contributing to ComfyUI-FlowControl!
|
||||
Thank you for your interest in contributing to **`ComfyUI-SaturnNodes`**! We welcome bug reports, feature suggestions, documentation updates, and pull requests from the community.
|
||||
|
||||
## Development Workflow
|
||||
Please take a moment to review these guidelines before submitting an issue or pull request.
|
||||
|
||||
1. Fork the repository on GitHub.
|
||||
2. Clone your fork locally into `ComfyUI/custom_nodes/ComfyUI-FlowControl`.
|
||||
3. Create a feature branch: `git checkout -b feature/my-feature`.
|
||||
4. Make your changes and test locally in ComfyUI.
|
||||
5. Ensure code is formatted cleanly and free of hardcoded personal paths or credentials.
|
||||
6. Commit your changes: `git commit -m "Add new feature"`.
|
||||
7. Push to your branch and submit a Pull Request.
|
||||
---
|
||||
|
||||
## Code Guidelines
|
||||
## 🐛 Reporting Issues & Bugs
|
||||
|
||||
- **Python**: Follow PEP 8 style guidelines. Ensure thread operations use atomic writes (`.tmp` + `os.replace`).
|
||||
- **JavaScript**: Follow ES Module conventions. Use `import { app } from "../../scripts/app.js"`.
|
||||
- **Security**: Never commit personal paths, hardcoded keys, or sensitive `.env` files.
|
||||
Before opening a new issue:
|
||||
1. **Check Existing Issues**: Search the [Issue Tracker](https://github.com/KOFiblto/ComfyUI-SaturnNodes/issues) to see if your bug or feature request has already been reported.
|
||||
2. **Include Reproduction Steps**: Provide a clear step-by-step description of how to reproduce the bug.
|
||||
3. **Include Logs & Environment Details**:
|
||||
- **ComfyUI Version** (V2 Vue UI or Classic V1 LiteGraph)
|
||||
- **OS** (Windows, Linux, macOS)
|
||||
- **Browser** (Chrome, Firefox, Edge)
|
||||
- Full Python console log traceback (if execution errored).
|
||||
|
||||
---
|
||||
|
||||
## 🛠️ Contribution Workflow (Feature Branches & PRs)
|
||||
|
||||
### Branch and Pull Request Policy
|
||||
Direct pushes to the `main` branch are strictly prohibited. Active development is organized into dedicated version branches (e.g. `v2.5.0`, `v2.6.0`). Community features or fixes should branch off the active version branch or submit Pull Requests targeting it.
|
||||
|
||||
Merging into `main` and creating releases is strictly gated by repository maintainer explicit authorization. Every release is tagged with `v` prefix (`vX.Y.Z`) and accompanied by release notes detailing all changes.
|
||||
|
||||
All development follows a **Branch $\rightarrow$ Pull Request** workflow:
|
||||
|
||||
```
|
||||
1. Fork/Branch Repository --> 2. Create Working Branch --> 3. Commit Changes --> 4. Open Pull Request
|
||||
```
|
||||
|
||||
### Step 1: Fork and Clone
|
||||
Fork the repository on GitHub, then clone your fork locally:
|
||||
```bash
|
||||
git clone https://github.com/YOUR_USERNAME/ComfyUI-SaturnNodes.git
|
||||
cd ComfyUI-SaturnNodes
|
||||
```
|
||||
|
||||
### Step 2: Create a Working Branch
|
||||
Create a descriptive branch for your change:
|
||||
```bash
|
||||
git checkout -b feature/short-description
|
||||
# Example: git checkout -b feature/live-preview-aspect-ratio
|
||||
```
|
||||
|
||||
### Step 3: Coding & Style Guidelines
|
||||
- **Python Style:** Match existing repository conventions. Keep imports at module scope. Do not add unnecessary `try/except` swallow blocks.
|
||||
- **No Internet Requests in Core Execution:** All node execution must run locally. Outbound network requests (e.g. Civitai/TMDB thumbnail scraping) must strictly require user consent via settings.
|
||||
- **No Hardcoded Keys or Absolute System Paths:** Do not commit personal system paths (`C:\Users\...` or `D:\...`), API keys, or private tokens.
|
||||
|
||||
### Step 4: Commit Your Changes
|
||||
Use concise, imperative subject lines for commits:
|
||||
- `Feat: ...` (New node or feature)
|
||||
- `Fix: ...` (Bug fix)
|
||||
- `Docs: ...` (Documentation update)
|
||||
- `Refactor: ...` (Code cleanup or restructuring)
|
||||
|
||||
Example:
|
||||
```bash
|
||||
git add .
|
||||
git commit -m "Fix: Resolve aspect ratio bounding box in Live Latent Preview"
|
||||
```
|
||||
|
||||
### Step 5: Push & Submit a Pull Request
|
||||
Push your branch to your fork on GitHub and open a Pull Request:
|
||||
```bash
|
||||
git push origin feature/short-description
|
||||
```
|
||||
1. Open [Pull Requests](https://github.com/KOFiblto/ComfyUI-SaturnNodes/pulls) $\rightarrow$ Click **New pull request**.
|
||||
2. Select the active version branch (e.g. `v2.5.0`) or `main` as base $\leftarrow$ `compare: feature/short-description`.
|
||||
3. Provide a brief summary of what changed and how you verified your fix.
|
||||
4. Wait for code review approval and automated status check validation!
|
||||
|
||||
---
|
||||
|
||||
## 📜 Code of Conduct & Mascot
|
||||
|
||||
- Be respectful and constructive when commenting on issues and reviewing pull requests.
|
||||
- The official mascot of ComfyUI is a very cute anime girl with massive fennec ears, a big fluffy tail, long blonde wavy hair, and blue eyes. Treat her with respect!
|
||||
|
||||
Thank you for helping make `ComfyUI-SaturnNodes` better for everyone!
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# AI Agent Guidelines
|
||||
|
||||
All AI agents (Claude, Gemini, Antigravity, Cursor, Copilot) contributing to this repository must follow the instructions in:
|
||||
- [`AGENTS.md`](./AGENTS.md) — Standard Operating Procedure & Architecture
|
||||
- [`CONTRIBUTING.md`](./CONTRIBUTING.md) — Git Workflow & Contribution Rules
|
||||
@@ -0,0 +1,271 @@
|
||||
# 🔍 Visual LoRA Loader & Batch Queue Logic Audit
|
||||
## Root Causes 1–5: Defect Analysis, Locations & Mechanics
|
||||
|
||||
> **Status:** Validated against `v2.6.3`. Root Causes **1** and **5** were genuine
|
||||
> defects and are now fixed. Root Cause **2** is intended persistent-cycle
|
||||
> behavior and is documented rather than changed. Root Causes **3** and **4** do
|
||||
> not reproduce with the current submission queue and were corrected below.
|
||||
|
||||
---
|
||||
|
||||
## Executive Overview of Observed Defects
|
||||
|
||||
When a user selects one or more LoRAs, queues a batch of 5, and while they are being queued selects different LoRA(s) and queues a batch of 10, the following symptoms occur:
|
||||
1. **LoRAs mixing together**: Prompts in Batch 2 receive LoRAs from Batch 1, or prompts in Batch 1 switch to LoRAs from Batch 2 mid-execution.
|
||||
2. **Duplicate generations**: A LoRA generates twice in a row when sequential single-generation was expected.
|
||||
3. **Missing LoRAs**: A LoRA that was explicitly selected is skipped and never generated.
|
||||
4. **Capture bypass**: The previous "workflow capture" (`Batch Graph Snapshot Guard`) failed to prevent these issues.
|
||||
|
||||
The audit identified **five candidate root causes**. Validation shows that the
|
||||
repeat/skip symptoms are explained by Root Causes **1** and **5**, while Root
|
||||
Cause **2** is by-design cycling; Root Causes **3** and **4** do not produce the
|
||||
described failures in the current code.
|
||||
|
||||
---
|
||||
|
||||
## 📌 Root Cause 1: Deterministic Hash Sorting vs. Unsorted Array Indexing
|
||||
|
||||
- **File**: [`nodes/lora_loader.py`](nodes/lora_loader.py)
|
||||
- **Functions**: `select_a_lora(...)`
|
||||
- **Verdict**: ✅ **Confirmed and fixed**
|
||||
- **Symptom Produced**: **Same LoRA generated twice; other selected LoRA skipped entirely.**
|
||||
|
||||
### What Is Happening
|
||||
In `select_a_lora`, the persistent state key on disk uses a SHA-256 hash calculated over the **alphabetically sorted** list of choices:
|
||||
|
||||
```python
|
||||
choices_hash = hashlib.sha256(
|
||||
json.dumps(sorted(valid_choices), ensure_ascii=False).encode("utf-8")
|
||||
).hexdigest()[:16]
|
||||
|
||||
state_key = f"{prefix}_{unique_id}_{choices_hash}"
|
||||
```
|
||||
|
||||
However, the selection indexes into the **unsorted** live list, whose order is
|
||||
the click/toggle order stored by the frontend (`web/js/folder_lora_loader.js`,
|
||||
`updateSelection`). Deselecting and reselecting a tile, or using a folder
|
||||
checkbox, changes that order without changing the hash.
|
||||
|
||||
### Failure Scenario:
|
||||
1. User selects **LoRA A**, then **LoRA B**:
|
||||
- `valid_choices` = `["A", "B"]`, `sorted(...)` = `["A", "B"]` → `choices_hash` = `H1`.
|
||||
- Prompt 1 executes at `index = 0` → **A**; saved `index = 1`.
|
||||
2. The user reorders the selection to **B**, then **A**:
|
||||
- `valid_choices` = `["B", "A"]`, sorted still `["A", "B"]` → same `H1`, so `index = 1` is loaded.
|
||||
- `valid_choices[1]` is **A** → **A generated twice, B skipped.**
|
||||
|
||||
### Fix Applied
|
||||
Sequential state now stores the order observed when the cycle started and
|
||||
indexes into that stored order. Reordering the selection can no longer move the
|
||||
index onto a different LoRA, so picks are never duplicated or skipped:
|
||||
|
||||
```python
|
||||
new_entry["index"] = index
|
||||
new_entry["order"] = order
|
||||
```
|
||||
|
||||
The index is also normalized (`index %= len(order)`), which keeps long-lived
|
||||
state bounded.
|
||||
|
||||
---
|
||||
|
||||
## 📌 Root Cause 2: Perpetual Inter-Batch Cycle Continuation
|
||||
|
||||
- **File**: [`nodes/lora_loader.py`](nodes/lora_loader.py)
|
||||
- **State**: `user/default/SaturnNodes/lora_loader_state.json`
|
||||
- **Verdict**: ✅ **Confirmed, but by design — documented, not changed**
|
||||
- **Symptom Produced**: **Batches starting on arbitrary mid-cycle LoRAs; leftover pools continuing across sessions.**
|
||||
|
||||
### What Is Happening
|
||||
`state_key` is keyed by `(mode, node_id, choices_hash)`; there is **no batch
|
||||
boundary, run ID, or queue session**. The `index` counter (Sequential) and the
|
||||
remaining `pool` (Random No-Replace) persist across batches, restarts, and days.
|
||||
Global state is pruned only by recency (100 most recent buckets).
|
||||
|
||||
### Failure Scenario:
|
||||
1. Three LoRAs (`[A, B, C]`), batch of 5 → A, B, C, A, B; saved `index = 5`.
|
||||
2. A later batch of 2 resumes at `5 % 3 = 2` → **C, A** instead of A, B.
|
||||
3. In **Random (No Replace)**, unpopped items from the previous run are consumed
|
||||
first, so a new batch can start on leftovers from the old pool.
|
||||
|
||||
### Why This Is Not Changed
|
||||
Continuous cycling is the documented "Sequential Cycle" / "Random (No Dupes)"
|
||||
contract, and `tests/test_lora_loader.py` asserts it across prompts. The suite
|
||||
already ships a manual `↺ Counter` reset and a per-node prompt→LoRA assignment
|
||||
map for retries. Resetting automatically per batch would silently change the
|
||||
documented cycling behavior, so this is treated as expected behavior with the
|
||||
reset control as the escape hatch.
|
||||
|
||||
---
|
||||
|
||||
## 📌 Root Cause 3: Premature Snapshot Guard Teardown in Frontend V2
|
||||
|
||||
- **File**: [`web/batch_queue.js`](web/batch_queue.js)
|
||||
- **Functions**: `setupQueueHooks()` (`app.queuePrompt` and `api.queuePrompt` hooks)
|
||||
- **Verdict**: ❌ **Not reproducible with the current submission queue**
|
||||
- **Original Claim**: A second batch makes ComfyUI's `app.queuePrompt` return
|
||||
`false` synchronously, so the wrapper's `finally` clears `activeBatchContext`
|
||||
before Batch 2's items are serialized, disabling the snapshot guard.
|
||||
|
||||
### Why the Claim Does Not Hold
|
||||
The wrapper defers the original call through `batchSubmissionQueue`:
|
||||
|
||||
```javascript
|
||||
const turn = batchSubmissionQueue.then(async () => {
|
||||
const batch = getNextBatch(count);
|
||||
activeBatchContext = batch;
|
||||
try {
|
||||
return await origQueuePrompt(number, batchCount, queueNodeIds);
|
||||
} finally {
|
||||
activeBatchContext = null;
|
||||
}
|
||||
});
|
||||
batchSubmissionQueue = turn.catch(() => {});
|
||||
return turn;
|
||||
```
|
||||
|
||||
`batchSubmissionQueue` is updated synchronously before returning, so a second
|
||||
queue action waits for the previous batch's submission promise to settle. Its
|
||||
callback — and therefore `origQueuePrompt` — is not invoked while
|
||||
`processingQueue === true`, so the synchronous `if (this.processingQueue)
|
||||
return false;` path cannot fire for the second UI batch. This serialization is
|
||||
present in the released `v2.6.2` as well.
|
||||
|
||||
ComfyUI's core behavior quoted in the original analysis is itself correct
|
||||
(`comfyui_frontend_package/static/assets`): `queuePrompt` pushes to
|
||||
`queueItems`, returns `false` while already processing, and otherwise loops with
|
||||
`await graphToPrompt()` per item. The defect was in assuming the wrapper reached
|
||||
that code immediately.
|
||||
|
||||
### Residual Edge Case
|
||||
The early teardown described above *can* occur if a non-SaturnNodes caller
|
||||
invokes `app.queuePrompt` while ComfyUI is already processing — the wrapper
|
||||
would then get the synchronous `false` and clear its context. That is outside
|
||||
the two-manual-batches scenario and is not addressed here.
|
||||
|
||||
---
|
||||
|
||||
## 📌 Root Cause 4: Single-Item Batches (`batchCount == 1`) Bypass Snapshot Guard
|
||||
|
||||
- **File**: [`web/batch_queue.js`](web/batch_queue.js)
|
||||
- **Functions**: `api.queuePrompt` hook (`isSnapshotGuardEnabled() && batchContext.batchCount > 1`)
|
||||
- **Verdict**: ⚠️ **Accurate code, not a defect**
|
||||
|
||||
### Why This Is Not a Defect
|
||||
The snapshot guard exists to keep items **inside one multi-item batch**
|
||||
consistent while the graph is re-serialized between iterations. A one-item batch
|
||||
has no later item to protect, so skipping the snapshot is harmless.
|
||||
|
||||
Additionally, the guard has never protected the *first* item of any batch: the
|
||||
snapshot is taken when that item is submitted, not when the user clicked Queue.
|
||||
Edits made between clicking and submission affect item 0 of a multi-item batch
|
||||
exactly as they affect a one-item batch. Tightening this would require capturing
|
||||
the prompt at click time, which is a separate design change, not a
|
||||
`batchCount > 1` bug.
|
||||
|
||||
---
|
||||
|
||||
## 📌 Root Cause 5: Function Signature Default Mismatch (`"All"` vs. `"Sequential"`)
|
||||
|
||||
- **File**: [`nodes/lora_loader.py`](nodes/lora_loader.py)
|
||||
- **Functions**: `VisualLoraLoader.INPUT_TYPES`, `VisualLoraLoader.IS_CHANGED`, `VisualLoraLoader.load_lora(...)`
|
||||
- **Verdict**: ✅ **Confirmed and fixed**
|
||||
- **Symptom Produced**: **Simultaneous loading of multiple LoRAs (stacking all models into one render).**
|
||||
|
||||
### What Is Happening
|
||||
`INPUT_TYPES` declared the hidden widget default as `"Sequential"`, while both
|
||||
`IS_CHANGED` and the execution method signature fell back to `"All"`:
|
||||
|
||||
```python
|
||||
# INPUT_TYPES
|
||||
"_selection_mode": ("STRING", {"default": "Sequential"}),
|
||||
|
||||
# IS_CHANGED / load_lora (before the fix)
|
||||
mode = kwargs.get("_selection_mode", "All")
|
||||
def load_lora(..., _selection_mode="All", ...)
|
||||
```
|
||||
|
||||
### Failure Scenario:
|
||||
1. A prompt arrives without `_selection_mode` (raw API call, third-party
|
||||
client, or hand-built prompt).
|
||||
2. The backend falls back to `"All"`, so the `share one selection` branch
|
||||
(`if _selection_mode in ["Random", "Sequential", "Random (No Replace)"]`) is
|
||||
skipped and `select_a_lora` is never called.
|
||||
3. Every selected LoRA is loaded onto the model and CLIP in one render instead
|
||||
of cycling one per prompt.
|
||||
|
||||
Normal UI use is unaffected: `web/js/folder_lora_loader.js` always creates a
|
||||
serialized `_selection_mode` widget defaulting to `"Sequential"`, which is why
|
||||
the bug only surfaced through API-style prompts. `IS_CHANGED` was also affected:
|
||||
its `"All"` fallback returned `""` (cacheable) for prompts missing the field.
|
||||
|
||||
### Fix Applied
|
||||
Both fallbacks now match the node definition and UI:
|
||||
|
||||
```python
|
||||
mode = kwargs.get("_selection_mode", "Sequential")
|
||||
def load_lora(..., _selection_mode="Sequential", ...)
|
||||
```
|
||||
|
||||
An explicit `"All"` sent by the UI still selects every LoRA simultaneously, so
|
||||
intentional "All Loras" mode is unchanged.
|
||||
|
||||
---
|
||||
|
||||
## Summary Matrix
|
||||
|
||||
| Root Cause | Affected File(s) | Primary Impact | Status |
|
||||
| :--- | :--- | :--- | :--- |
|
||||
| **1. Hash vs Indexing Order** | `nodes/lora_loader.py` | Same LoRA selected twice; other selected LoRAs skipped entirely. | ✅ Fixed (stored cycle order) |
|
||||
| **2. Inter-Batch State Carryover** | `nodes/lora_loader.py` | Prompts start at arbitrary cycle offsets; leftover pools continue. | ✅ By design (manual `↺ Counter` reset) |
|
||||
| **3. Snapshot Guard Teardown** | `web/batch_queue.js` | — | ❌ Not reproducible; submission queue serializes batches |
|
||||
| **4. Batch Count == 1 Bypass** | `web/batch_queue.js` | — | ⚠️ Expected; no second item to snapshot |
|
||||
| **5. Default Mode Mismatch** | `nodes/lora_loader.py` | API prompts silently stack all selected LoRAs. | ✅ Fixed (defaults aligned to `Sequential`) |
|
||||
|
||||
### Regression Coverage Added
|
||||
|
||||
- `test_sequential_reorder_keeps_cycle_position` — reordering the selection
|
||||
continues the existing cycle without repeating or skipping a LoRA.
|
||||
- `test_sequential_changed_selection_starts_fresh_cycle` — changing the selected
|
||||
set still starts a new bucket.
|
||||
- `test_missing_selection_mode_defaults_to_sequential` — `INPUT_TYPES`,
|
||||
`load_lora`, and `IS_CHANGED` agree on the fallback mode.
|
||||
|
||||
---
|
||||
|
||||
## Follow-up: Click-to-Queue Bleed (reported after the initial audit)
|
||||
|
||||
**Symptom:** Queue a batch with LoRA A, then queue more prompts with LoRAs B/C
|
||||
10× and A again 5× while the earlier batch is still loading, and prompts come
|
||||
out with text or LoRA selections taken from a different click ("character
|
||||
bleeding"). Rapid Queue clicks made the effect worse.
|
||||
|
||||
**Actual cause (client-side timing, not parallelism):** ComfyUI's
|
||||
`app.queuePrompt` does not serialize a click immediately when it is already
|
||||
submitting another batch. It pushes the request onto `queueItems` and returns
|
||||
`false`; the graph is serialized with `graphToPrompt()` only when that request's
|
||||
turn starts. SaturnNodes additionally serializes UI queue actions through
|
||||
`batchSubmissionQueue`. Any canvas edit made between the click and the deferred
|
||||
serialization was therefore captured by the queued prompt. The old snapshot
|
||||
guard only froze items **after** item 0 of a multi-item batch, so item 0 and
|
||||
every single-item queue remained exposed.
|
||||
|
||||
**Fix:** `web/batch_queue.js` now captures the graph at click time whenever a
|
||||
submission is already pending (`pendingSubmissions > 0` or
|
||||
`app.processingQueue === true`). Item 0 of that batch uses the captured prompt,
|
||||
later items reuse it as the snapshot, and `mergeDynamicSeeds` keeps randomized
|
||||
seeds from the live prompt so generated seeds still match the UI. Idle clicks
|
||||
skip the extra serialization.
|
||||
|
||||
**Prompt iterator:** reviewed for the same class of bug. It is not affected by
|
||||
submission timing because ComfyUI executes prompts serially on the server. One
|
||||
real bleed source was fixed: the optional `prompt` input was ignored whenever
|
||||
the node's local `text` widget was non-empty, so stale local characters could
|
||||
be replayed instead of the connected *Prompt Counter* output. A connected input
|
||||
now always wins.
|
||||
|
||||
**Why the old tests missed it:** the Python tests call selection/iterator
|
||||
functions directly and the DOM tests never modelled a pending HTTP submission.
|
||||
The new `tests/test_batch_queue_snapshot_dom.mjs` reproduces the deferred
|
||||
`graphToPrompt()` order of ComfyUI, and the iterator precedence tests live in
|
||||
`tests/test_prompt_iterator.py`.
|
||||
@@ -1,133 +1,522 @@
|
||||
# 🍃 ComfyUI-FlowControl
|
||||
# 🪐 ComfyUI-SaturnNodes
|
||||
|
||||
A unified custom node suite for **ComfyUI** featuring real-time queue controls, mid-generation pausing, queue crash recovery, visual model pickers, live latent canvas previews, folder automation, and prompt utilities.
|
||||
A unified workflow control and productivity suite for **ComfyUI**. Both **Frontend V1 (LiteGraph)** and **Frontend V2 (Vue UI)** are supported, though `VisualLoraLoader` breaks on V1 for now (will be fixed).
|
||||
|
||||
> ⚠️ **Important Notice**: This custom node pack (`ComfyUI-FlowControl`) is engineered exclusively for the **New ComfyUI Frontend (v2)**. The legacy/old ComfyUI frontend is not supported.
|
||||
> 📖 **New here?** Follow the step-by-step [Complete Walkthrough](WALKTHROUGH.md) for setup, node recipes, and batch-prompting patterns.
|
||||
|
||||
---
|
||||
|
||||
## 🌟 Features & Included Nodes
|
||||
## 🚀 Key Capabilities & Overview
|
||||
|
||||
Every node in this suite is prefixed with the signature **`🍃`** brand icon for quick identification on your canvas and search menu.
|
||||
|
||||
### ⏸️ Queue & Workflow Control (`FlowControl/Queue`)
|
||||
|
||||
1. **`🍃 Pause Queue` (Global Integration)**
|
||||
- Adds a native-styled **Pause / Continue** button directly into the ComfyUI V2 top action bar.
|
||||
- **Pause (Finish)**: Completes the currently active generation, then holds the queue before starting the next prompt.
|
||||
- **Pause (Instant)**: Pauses immediately before the next node executes in the active workflow (resumable mid-workflow without losing progress).
|
||||
- **Visual Indicators**: Green when unpaused, Orange when pausing/paused with pulsing glow animation.
|
||||
- **Status Text**: Displays `Pausing (Finish)...` or `Pausing (Instant)...` while waiting for execution to reach the pause gate, then transforms into `Continue` with a play icon `▶`.
|
||||
|
||||
2. **`🍃 Persistent Queue` (Global Integration)**
|
||||
- Automatically saves all queued workflows to `persistent_queue.json` in real time on disk.
|
||||
- If ComfyUI crashes, suffers a power outage, or is closed, all unfinished prompts are automatically restored to the queue on next boot in a paused state.
|
||||
- **Session Ownership Claiming**: Automatically re-assigns restored queue items to your active browser tab's `client_id` when you open the page, removing the *"Running in another tab"* warning.
|
||||
|
||||
3. **`🍃 FlowControl Decision`**
|
||||
- A drop-in node that pauses workflow execution at a specific point and prompts you with a UI popup.
|
||||
- **Continue**: Resumes execution normally (outputs `False`).
|
||||
- **Cancel**: Cancels this branch of the workflow (outputs `True`, allowing you to route it into a Switch node).
|
||||
- **Stop Workflow**: Instantly kills the entire active generation queue.
|
||||
- **Windows Notification**: Optionally sends a native Windows 11 toast notification when it's waiting for your input.
|
||||
- **Live Latent Preview**: Watch generations materialize step-by-step directly inside your graph canvas nodes in real time.
|
||||
- **Stereoscopic 3D & Cross-Eye Generator**: Convert 2D artwork and renders into 3D stereograms, parallel views, and red-cyan anaglyphs with automatic Depth Anything V2 fallback.
|
||||
- **Batch Prompting & Iteration**: Queue multiple prompts automatically with delimiter splits, regex blocks, and persistent queue state tracking.
|
||||
- **Visual Thumbnail Browsing**: Browse LoRAs and image folders with auto-scraped Civitai & TMDB previews, ranking badges, and multi-selection.
|
||||
- **Real-Time Queue & Pause Controls**: Pause, resume, or cleanly cancel execution mid-generation, with automatic crash recovery for interrupted queues.
|
||||
- **Aspect Ratio & Resolution Utilities**: Find exact matching aspect ratios, crop dimensions, and visual sizing helpers.
|
||||
- **Safe Local Script Execution**: Run local scripts and executables strictly confined to `ComfyUI/scripts/` with interactive operator authorization.
|
||||
|
||||
---
|
||||
|
||||
### 🎨 Visual Loaders & Pickers (`FlowControl/Loaders`)
|
||||
## 📦 Individual Nodes Reference
|
||||
|
||||
3. **`🍃 Visual LoRA Loader`**
|
||||
- Visual LoRA selector supporting folder filtering (e.g. `Celebrity/*`) and formatted pretty names (e.g. `Ana De Armas V1`).
|
||||
- Automatically tracks usage frequency (`lora_usage.json`).
|
||||
- **Custom Preview Images**: Simply save a preview image file with the exact same name as your LoRA in the same folder (e.g. `my_lora.png`, `my_lora.jpg`, or `my_lora.webp` next to `my_lora.safetensors`).
|
||||
- **Automated Civitai SHA256 Search**: Computes SHA256 file hashes and queries Civitai's API (`/api/v1/model-versions/by-hash/{hash}`) with optional `CIVITAI_API_KEY` authorization to download model preview thumbnails.
|
||||
- **TMDB Fallback**: Optional celebrity preview search via TMDB API key.
|
||||
Expand any node below to view its description, inputs, outputs, and usage documentation.
|
||||
|
||||
4. **`🍃 Visual Image Loader`**
|
||||
- Visual image folder browser with instant visual preview pickers.
|
||||
- Extracts positive prompt metadata from PNG parameters and EXIF user comments.
|
||||
- Outputs `IMAGE`, `positive_prompt`, `width`, and `height`.
|
||||
<details>
|
||||
<summary><b>🪐 📁 LoRA Loader (Folder)</b> (<code>FolderLoraLoader</code>)</summary>
|
||||
|
||||
5. **`🍃 Recent Outputs`**
|
||||
- Cycles through recently generated images by step index for seamless integration with ComfyUI's sidebar/assets panel.
|
||||
- Supports custom `output_folder` paths.
|
||||
- Uses an `index` paired with an `amount` to chronologically step through the X newest images in a batch.
|
||||
#### Overview
|
||||
Loads a LoRA by folder path using raw filename matching or formatted names with customizable output formatting options.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`model`** (`MODEL`): Input model.
|
||||
- **`clip`** (`CLIP`): Input CLIP.
|
||||
- **`folder`** (`STRING`): Subfolder filter (e.g. `Characters`, `Styles/Anime`, `*`).
|
||||
- **`lora_name`** (`COMBO`): Dropdown listing available LoRA files in the target folder.
|
||||
- **`strength_model`** (`FLOAT`): Model weight strength (default `1.0`).
|
||||
- **`strength_clip`** (`FLOAT`): CLIP weight strength (default `1.0`).
|
||||
- **`output_format`** (`COMBO`, *Advanced*): Choose output format (`Parsed Name`, `Filename`, `Filename without extension`, `Relative Path`, `Full Path`, `Custom Regex`).
|
||||
- **`custom_regex`** (`STRING`, *Advanced*): Pattern used when `output_format` is `Custom Regex`.
|
||||
|
||||
#### Outputs
|
||||
- **`MODEL`**: Patched model.
|
||||
- **`CLIP`**: Patched CLIP.
|
||||
- **`lora_name`** (`STRING`): Formatted name or filename of the loaded LoRA.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 ✨ LoRA Loader (Pretty)</b> (<code>FolderLoraLoaderPretty</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Loads a LoRA using formatted pretty names (e.g. `Cyberpunk Character V1` instead of `char_Cyberpunk-Character_v1.safetensors`).
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`model`** (`MODEL`): Input model.
|
||||
- **`clip`** (`CLIP`): Input CLIP.
|
||||
- **`folder`** (`STRING`): Subfolder filter.
|
||||
- **`lora_name`** (`COMBO`): Dropdown listing formatted LoRA names.
|
||||
- **`strength_model`** (`FLOAT`): Model weight strength.
|
||||
- **`strength_clip`** (`FLOAT`): CLIP weight strength.
|
||||
|
||||
#### Outputs
|
||||
- **`MODEL`**: Patched model.
|
||||
- **`CLIP`**: Patched CLIP.
|
||||
- **`lora_name`** (`STRING`): Selected LoRA name.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 🖼️ Visual LoRA Loader</b> (<code>VisualLoraLoader</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Visual thumbnail browser for LoRAs with Civitai SHA256 search & TMDB auto-scraping, popularity rank badges (🔥), multi-selection support, and independent per-node thumbnail zoom sizing.
|
||||
|
||||
#### Maintenance Reset Buttons
|
||||
The node's control bar exposes three maintenance buttons — each asks for confirmation before acting:
|
||||
- **`↺ Counter`**: Resets this node's `Sequential`/`Random` cycle position and clears its prompt→LoRA assignments (stored in `user/default/SaturnNodes/lora_loader_state.json`). ⚠️ Only use between batches: running/queued workflows will restart the cycle from the first selected LoRA and can mismatch their prompts.
|
||||
- **`🧹 Usage`**: Clears all LoRA usage counters and badges (`lora_usage.json`). Cosmetic only — affects badges and `Usage` sortings, not LoRA loading.
|
||||
- **`🗑 Cache`**: Deletes cached thumbnail images. They regenerate automatically when you view the grid.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`model`** (`MODEL`): Input model.
|
||||
- **`clip`** (`CLIP`): Input CLIP.
|
||||
- **`folder`** (`STRING`): Subfolder filter.
|
||||
- **`strength_model`** (`FLOAT`): Model strength.
|
||||
- **`strength_clip`** (`FLOAT`): CLIP strength.
|
||||
- **`tile_size`** (`SLIDER`, UI): Adjust thumbnail card size (persisted per node in workflow JSON).
|
||||
- **`display_mode`** (`COMBO`, *Advanced*): `Scrollable` vs `Show All`.
|
||||
- **`sort_loras_by`** (`COMBO`, *Advanced*): Sort by `Name (A-Z)`, `Name (Z-A)`, `Date Modified (Newest First)`, `Date Modified (Oldest First)`, or `Usage (Most Used First)`.
|
||||
|
||||
#### Outputs
|
||||
- **`MODEL`**: Patched model.
|
||||
- **`CLIP`**: Patched CLIP.
|
||||
- **`lora_name`** (`STRING`): Comma-separated list of active LoRA names.
|
||||
|
||||
> [!NOTE]
|
||||
> Frontend V1 and V2 are supported, though `VisualLoraLoader` breaks on Frontend V1 for now (will be fixed). Frontend V2 (Vue UI) is recommended for visual LoRA browsing.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 📷 Visual Image Loader</b> (<code>VisualImageLoader</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Visual thumbnail browser for image folders with instant preview selection, independent per-node thumbnail zoom sizing, and EXIF positive prompt metadata extraction. Path access is strictly confined to ComfyUI `input`, `output`, and `temp` directories for safe operation.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`folder`** (`STRING`): Folder path to load images from (confined to ComfyUI input/output/temp).
|
||||
- **`tile_size`** (`SLIDER`, UI): Adjust thumbnail card size (persisted per node in workflow JSON).
|
||||
- **`display_mode`** (`COMBO`, *Advanced*): `Scrollable` vs `Show All`.
|
||||
- **`sort_images_by`** (`COMBO`, *Advanced*): `Name (A-Z)`, `Name (Z-A)`, `Date Modified (Newest First)`, `Date Modified (Oldest First)`.
|
||||
|
||||
Folder scans run off ComfyUI's event loop, thumbnails are cached on disk, and prompt metadata is extracted only for the newest `IMAGE_LOADER_METADATA_LIMIT` images (default `128`, configurable in Settings).
|
||||
|
||||
#### Outputs
|
||||
- **`IMAGE`**: Selected image tensor.
|
||||
- **`positive_prompt`** (`STRING`): Extracted prompt metadata.
|
||||
- **`width`** (`INT`): Image width.
|
||||
- **`height`** (`INT`): Image height.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 📂 Load Image From Folder</b> (<code>LoadImageFromFolder</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Monitors a folder for incoming images, loads the target image into a PyTorch tensor, with an optional toggle to delete the image after loading. Path access and deletion are strictly confined to ComfyUI `input` and `output` directories.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`folder`** (`STRING`): Watch directory path (default `input/watch`).
|
||||
- **`wait_if_folder_is_empty`** (`BOOLEAN`): Poll until image arrives vs return dummy tensor immediately.
|
||||
- **`rescan_interval`** (`INT`): Polling frequency in seconds.
|
||||
- **`sort_by`** (`COMBO`): Sort by `date_modified`, `date_created`, or `name`.
|
||||
- **`regex_filter`** (`STRING`): Regex pattern to filter filenames.
|
||||
- **`delete_image`** (`BOOLEAN`, *Default: false*): If enabled, deletes the image file from disk after loading (strictly inside `input` or `output`).
|
||||
|
||||
#### Outputs
|
||||
- **`image`** (`IMAGE`): Loaded image tensor.
|
||||
- **`has_image`** (`BOOLEAN`): True if image loaded, False if folder empty.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 ⏱️ Recent Outputs</b> (<code>LoadRecentOutputs</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Loads the N newest images from an output directory with step-through index selection.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`output_folder`** (`STRING`): Target output directory.
|
||||
- **`amount`** (`INT`): Number of newest images to pool (1-100).
|
||||
- **`index`** (`INT`): Step index selection.
|
||||
|
||||
#### Outputs
|
||||
- **`IMAGE`**: Output image tensor.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 👁️ Live Latent Preview</b> (<code>PreviewLatentLiveNode</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Canvas rendering node that listens to sampler WebSocket latent binary streams and displays real-time live previews on the canvas during generation.
|
||||
|
||||
#### Inputs & Outputs
|
||||
- **Category**: `🪐 SaturnNodes/Previews`
|
||||
- **Output Node**: True
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 ⏸️ Saturn Decision</b> (<code>SaturnDecision</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Pauses workflow execution at a specific step and displays an inline UI popup with Continue, Cancel, or Stop Workflow actions, plus optional native OS desktop notifications.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`disable`** (`BOOLEAN`): Bypass decision gate.
|
||||
- **`send_os_notification`** (`BOOLEAN`): Trigger native OS desktop toast on pause.
|
||||
- **`timeout`** (`INT`): Auto-continue timeout in seconds (-1 = infinite).
|
||||
|
||||
#### Outputs
|
||||
- **`cancel`** (`BOOLEAN`): False on Continue, True on Cancel (for branch routing).
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 📐 Text Aspect Ratio Finder</b> (<code>TextAspectRatioFinder</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Parses input text for aspect ratios (e.g. `16:9`, `2.35:1`), syntax-checks them, and calculates pixel resolution for a target megapixel target.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`aspect_ratios`** (`STRING`): Allowed ratio list (e.g. `16:9, 9:16, 1:1`). When populated, strictly enforces matching only the configured ratios and ignores any other numbers/ratios in text. When left completely empty (`""`), acts in open mode and accepts any valid ratio found in text.
|
||||
- **`search_mode`** (`COMBO`): `First match (Front)` vs `Last match (Back)`.
|
||||
- **`target_mp`** (`FLOAT`): Target megapixels (e.g. `1.0 MP`).
|
||||
- **`default_aspect_ratio`** (`COMBO`): Fallback ratio if none found.
|
||||
- **`multiple_of`** (`INT`, *Advanced*): Dimension step multiple (default `8`).
|
||||
- **`min_mp`** (`FLOAT`, *Advanced*): Minimum megapixel limit.
|
||||
- **`max_mp`** (`FLOAT`, *Advanced*): Maximum megapixel limit.
|
||||
- **`text`** (`STRING`, *Optional Input*): Input text to search.
|
||||
|
||||
#### Outputs
|
||||
- **`width`** (`INT`): Computed width in pixels.
|
||||
- **`height`** (`INT`): Computed height in pixels.
|
||||
- **`aspect_ratio`** (`STRING`): Detected or fallback aspect ratio string.
|
||||
- **`cleaned_text`** (`STRING`): Input text with the aspect ratio token stripped.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 📐 Preview Image Size & Aspect Ratio</b> (<code>PreviewImageSizeAspectRatio</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Computes image dimensions and aspect ratio from an input image tensor and formats the result for previewing.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`image`** (`IMAGE`): Input image tensor.
|
||||
|
||||
#### Outputs
|
||||
- **`width`** (`INT`): Image width.
|
||||
- **`height`** (`INT`): Image height.
|
||||
- **`aspect_ratio`** (`STRING`): Closest matched aspect ratio string.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 🔎 Text LoRA Finder & Loader</b> (<code>TextLoraFinder</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Scans input prompt text for `<lora:name:strength>` tags or names matching a folder on disk, dynamically applies them, and returns patched models along with sanitized prompt text.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`model`** (`MODEL`): Input model.
|
||||
- **`clip`** (`CLIP`): Input CLIP.
|
||||
- **`folder`** (`STRING`): Subfolder to search for LoRA files.
|
||||
- **`fallback_strength_model`** (`FLOAT`): Default model strength if omitted in prompt tag.
|
||||
- **`fallback_strength_clip`** (`FLOAT`): Default CLIP strength if omitted in prompt tag.
|
||||
- **`clean_prompt`** (`BOOLEAN`): Strips matched `<lora:...>` tags from output text.
|
||||
- **`text`** (`STRING`, *Optional Input*): Input prompt text.
|
||||
|
||||
#### Outputs
|
||||
- **`MODEL`**: Patched model.
|
||||
- **`CLIP`**: Patched CLIP.
|
||||
- **`text`** (`STRING`): Cleaned prompt text.
|
||||
- **`lora_name`** (`STRING`): Comma-separated list of loaded LoRAs.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 🔄 Prompt Queue Iterator</b> (<code>PromptQueueIterator</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Deterministically iterates over multiline prompt text blocks per queue run with live progress display, counter reset controls, and index tracking. Supports splitting prompts by empty lines, newlines, or a custom regular expression delimiter.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`pop_mode`** (`COMBO`): `Sequential (Loop on End)`, `Sequential (Stop on End)`, `Random (Keep)`, `Random (Cycle)`.
|
||||
- **`track_by`** (`COMBO`): `Text & Node` (default; the pointer is isolated per prompt text and splitter configuration) or `Node Only` (one pointer per graph node that survives prompt edits and splitter changes).
|
||||
- **`separator`** (`COMBO`): Split prompts by `>1 Empty Line`, `Newline`, `>2 Empty Lines`, or `Custom Regex`.
|
||||
- **`custom_regex`** (`STRING`): Custom regular expression delimiter (e.g. `\n---\n`). Automatically enabled when `separator` is set to `Custom Regex` and disabled otherwise.
|
||||
- **`text`** (`STRING`, Multiline): Multiline prompt text containing your queued blocks.
|
||||
- **`prompt`** (`STRING`, Optional Input): Connect an external multiline string node (such as *Prompt Counter*).
|
||||
|
||||
#### Outputs
|
||||
- **`prompt`** (`STRING`): The selected prompt block for the current queue run.
|
||||
- **`remaining_text`** (`STRING`): All remaining prompt blocks joined by the delimiter. For `Random (Keep)` this is empty because random selection does not consume an ordered queue.
|
||||
- **`remaining_count`** (`INT`): Count of remaining items in the current cycle (always `0` for `Random (Keep)`).
|
||||
- **`exhausted`** (`BOOLEAN`): `True` when `Sequential (Stop on End)` has reached the final prompt (subsequent runs repeat the last prompt and remain exhausted).
|
||||
|
||||
Clicking `Reset Counter (0)` now also clears crash-recovery prompt assignments and random shuffle permutations for the node, so the next run starts from a clean state.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 📝 Prompt Counter</b> (<code>PromptCounter</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Multiline text prompt input node that functions identically to a standard multiline string node, with real-time prompt counting according to your chosen delimiter (`>1 Empty Line`, `Newline`, `>2 Empty Lines`, or `Custom Regex`). Displays the live prompt count directly on the node on every keystroke, allowing you to easily determine and set the exact ComfyUI batch size before running your queue.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`text`** (`STRING`, Multiline): Input prompt text block.
|
||||
- **`separator`** (`COMBO`, *Advanced*): Prompt delimiter (`>1 Empty Line`, `Newline`, `>2 Empty Lines`, `Custom Regex`). Hidden by default under Advanced Options.
|
||||
- **`custom_regex`** (`STRING`, *Advanced*): Custom regex delimiter pattern (e.g. `\n---\n`). Dynamically enabled/disabled based on delimiter selection.
|
||||
|
||||
#### Outputs
|
||||
- **`STRING`**: The unmodified prompt text (ready to connect into samplers or *Prompt Queue Iterator*).
|
||||
- **`count`** (`INT`): The registered prompt count integer.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 🔤 Multi Text Replacer</b> (<code>MultiTextReplacer</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Performs multiple text replacements in a single step using comma-separated or newline-separated find/replace lists, with support for exact phrases and regex matching.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`text`** (`STRING`, Multiline): Input text to modify.
|
||||
- **`find_text`** (`STRING`, Multiline): Search terms (comma-separated or lines).
|
||||
- **`replace_text`** (`STRING`, Multiline): Replacement terms.
|
||||
- **`case_sensitive`** (`BOOLEAN`): Toggle case-sensitive matching.
|
||||
- **`use_regex`** (`BOOLEAN`): Treat find patterns as regular expressions.
|
||||
|
||||
#### Outputs
|
||||
- **`text`** (`STRING`): Modified text.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 ✂️ Text Split</b> (<code>SaturnTextSplit</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Splits text into two parts at a specified delimiter. Supports forward (from start) and backward (from end) search, as well as regular expressions.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`text`** (`STRING`, Multiline): Input text to split.
|
||||
- **`split_by`** (`STRING`): Character sequence or regex pattern to split on (e.g. `--`, `,`, `---`, `\n`).
|
||||
- **`split_direction`** (`COMBO`, *Advanced*): `forward (first occurrence from start)` vs `backward (last occurrence from end)`.
|
||||
- **`use_regex`** (`BOOLEAN`, *Advanced*): Treat delimiter as a regular expression pattern.
|
||||
- **`strip_whitespace`** (`BOOLEAN`, *Advanced*): Trims surrounding whitespace from output strings.
|
||||
|
||||
#### Outputs
|
||||
- **`text1`** (`STRING`): Text before the split delimiter.
|
||||
- **`text2`** (`STRING`): Text after the split delimiter.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 ⚡ Run Local File</b> (<code>RunLocalFileNode</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Safely executes a local script or executable (`.bat`, `.cmd`, `.ps1`, `.exe` on Windows; `.sh` or binaries on Linux/macOS) with command-line parameters and working directory control. Built with multi-layered security guards to ensure malicious workflows or dropped images cannot execute code without explicit authorization.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> **Scripts Directory Restriction**: For operator security, this node can **ONLY execute scripts and binaries located inside the `ComfyUI/scripts/` directory** (e.g. `ComfyUI/scripts/my_script.bat`). Relative subfolders (such as `subfolder/process.sh`) are supported as long as they resolve inside `ComfyUI/scripts/`. Absolute paths (e.g. `C:\...` or `/bin/...`) and directory traversal sequences (`..`) are permanently blocked.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`file_path`** (`STRING`): Relative path to the executable or script (e.g. `process.bat`, `subfolder/script.sh`). Strictly confined to `ComfyUI/scripts/` with absolute path and traversal (`..`) rejection.
|
||||
- **`parameters`** (`STRING`, Multiline, Optional): Command-line arguments. Supports spaces and standard quoting (parsed safely via `shlex` without shell expansion).
|
||||
- **`working_directory`** (`STRING`, Optional): Execution working directory (`cwd`). Defaults to the folder containing the script if left blank.
|
||||
- **`run_mode`** (`COMBO`): `Synchronous (Wait for Output)` vs `Asynchronous (Background)`.
|
||||
- **`timeout`** (`INT`): Maximum execution time in seconds for synchronous mode (`0` = no timeout). If exceeded, the process tree is cleanly terminated via OS process management.
|
||||
- **`trigger`** (`*`, Optional Wildcard): Input wire allowing you to connect any workflow data to order execution.
|
||||
- **⚡ In-Memory Authorization**: Does NOT use an input widget (preventing malicious authorization smuggling inside shared workflow JSON or image metadata). Instead, the operator must click the interactive `⚡ Authorize Run (5 min)` button directly on the canvas node to grant a 5-minute single-use execution token.
|
||||
|
||||
#### Outputs
|
||||
- **`stdout`** (`STRING`): Captured standard output stream.
|
||||
- **`stderr`** (`STRING`): Captured standard error stream.
|
||||
- **`exit_code`** (`INT`): Process exit return code (`0` = success, `-1` = blocked / timed out).
|
||||
- **`success`** (`BOOLEAN`): `True` if `exit_code == 0`, `False` otherwise.
|
||||
- **`passthrough`** (`*`): Passes input `trigger` data through untouched to downstream nodes.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🪐 👓 Stereoscopic 3D / Cross-Eye Generator</b> (<code>SaturnStereo3D</code>)</summary>
|
||||
|
||||
#### Overview
|
||||
Inspired by Stereonix, this node transforms any 2D image into stereoscopic 3D pairs for cross-eye viewing, parallel 3D viewers, VR/3D displays, and red-cyan anaglyph 3D glasses using depth-based horizontal disparity mapping (DIBR). Features automatic on-demand Depth Anything V2 Small fallback, crop/zoom/pan controls, outer border framing, empty gap spacing, and focus alignment guide dots.
|
||||
|
||||
#### Inputs & Widgets
|
||||
- **`image`** (`IMAGE`): The source 2D image.
|
||||
- **`mode`** (`COMBO`): 3D viewing format (`Cross-Eye`, `Parallel (Wall-Eyed)`, `Red-Cyan Anaglyph`, `Side-by-Side (Parallel)`, `Top-Bottom (Over-Under)`, `Left Eye Only`, `Right Eye Only`).
|
||||
- **`depth_intensity`** (`FLOAT` slider, `0.0` - `100.0`, default `25.0`): Controls how much the left and right eye perspectives shift apart horizontally.
|
||||
- **`depth_gamma`** (`FLOAT` slider, `0.1` - `3.0`, default `0.50`): Adjusts the depth curve for realistic falloff.
|
||||
- **`invert_depth`** (`BOOLEAN`, default `False`): Flips near and far depth planes.
|
||||
- **`convergence_plane`** (`FLOAT` slider, `0.0` - `1.0`, default `0.50`): Focal depth plane where horizontal shift is zero (the screen surface). Depths closer pop out; depths farther sink in.
|
||||
- **`image_zoom`** (`FLOAT` slider, `10.0` - `300.0`, default `100.0`): Zoom into areas of interest before generating 3D disparity.
|
||||
- **`pan_horizontal`** (`FLOAT` slider, `-100.0` - `100.0`, default `0.0`): Horizontal pan offset.
|
||||
- **`pan_vertical`** (`FLOAT` slider, `-100.0` - `100.0`, default `0.0`): Vertical pan offset.
|
||||
- **`gap_spacing`** (`INT`, default `20`): Empty black space between the left and right eye panels.
|
||||
- **`outer_border`** (`INT`, default `0`): Black border padding around the canvas (allows shrinking images for easier cross-eye viewing on large monitors).
|
||||
- **`alignment_dots`** (`BOOLEAN`, default `True`): Renders focus guide dots centered above each eye view to easily lock into 3D cross-eye view.
|
||||
- **`fill_method`** (`COMBO`): Hole-filling method (`Smooth (Bilinear Grid)` vs `Nearest Neighbor (Stereonix)`).
|
||||
- **`preview_zoom`** (`COMBO`): Scale preview for comfortable canvas cross-eye viewing (`Fit`, `25%`, `33%`, `50%`, `75%`, `100%`).
|
||||
- **`depth_map`** (`IMAGE`, *Optional Input*): Optional depth map from any node. If not connected, automatically downloads and runs Depth Anything V2 Small ONNX model internally.
|
||||
|
||||
#### Outputs
|
||||
- **`stereogram`** (`IMAGE`): The composited 3D stereogram image.
|
||||
- **`left_eye`** (`IMAGE`): The warped left eye viewpoint.
|
||||
- **`right_eye`** (`IMAGE`): The warped right eye viewpoint.
|
||||
- **`depth_map`** (`IMAGE`): The normalized depth map used for disparity calculation.
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
### ⚙️ Automation, Previews & Utilities (`FlowControl/Automation`, `FlowControl/Previews`, `FlowControl/Utils`)
|
||||
|
||||
6. **`🍃 Load image from folder`**
|
||||
- Monitors a directory for incoming images (`.png`, `.jpg`, `.jpeg`, `.webp`), converts them into PyTorch image tensors, and removes processed files.
|
||||
- **`wait_if_folder_is_empty` Toggle**:
|
||||
- `True` (Default): Safely polls until an image arrives in the directory at your chosen `rescan_interval` (seconds).
|
||||
- `False`: Immediately checks the directory. If an image is present, loads it. If no image is present, returns a dummy tensor and `has_image=False` immediately without blocking.
|
||||
- **Advanced Sorting**: Supports sorting incoming images by `date_modified`, `date_created`, or `name`.
|
||||
- **Regex Filtering**: Use standard Python regex strings (e.g. `^img_.*\.png$`) to target specific file formats or naming conventions.
|
||||
|
||||
7. **`🍃 Undo Placeholder`**
|
||||
- Scans text prompt strings for names matching LoRA files in a specified folder filter and restores placeholder tokens like `%celeb%`.
|
||||
|
||||
8. **`🍃 Live Latent Preview`**
|
||||
- Custom canvas node that listens to WebSocket latent preview binary streams (`b_preview`, `b_preview_with_metadata`) during sampler execution and renders real-time sampler previews directly inside the node canvas.
|
||||
|
||||
---
|
||||
|
||||
## 📋 Complete Node Overview Table
|
||||
|
||||
| Node Class | Display Name | Category | Primary Function |
|
||||
| :--- | :--- | :--- | :--- |
|
||||
| `FlowControlDecision` | **`🍃 FlowControl Decision`** | `FlowControl/Queue` | Inline workflow pausing with Continue/Cancel UI and native notifications |
|
||||
| `FolderLoraLoaderVisualPrettyV2` | **`🍃 Visual LoRA Loader`** | `FlowControl/Loaders` | Visual LoRA picker with SHA256 Civitai/TMDB previews |
|
||||
| `ImageLoaderVisualPrettyV2` | **`🍃 Visual Image Loader`** | `FlowControl/Loaders` | Visual image browser with EXIF prompt extraction |
|
||||
| `LoadRecentOutputs` | **`🍃 Recent Outputs`** | `FlowControl/Loaders` | Cycles recent output images for assets/sidebar |
|
||||
| `LoadImageFromFolder` | **`🍃 Load image from folder`** | `FlowControl/Automation` | Folder image watcher with regex and intelligent sorting |
|
||||
| `UndoPlaceholder` | **`🍃 Undo Placeholder`** | `FlowControl/Utils` | Replaces pretty LoRA names with `%placeholder%` tokens |
|
||||
| `PreviewLatentLive` | **`🍃 Live Latent Preview`** | `FlowControl/Previews` | Canvas real-time sampler latent preview display |
|
||||
|
||||
---
|
||||
|
||||
## ⚙️ ComfyUI Settings Menu Integration
|
||||
<a id="comfyui-settings-menu-reference"></a>
|
||||
## ⚙️ ComfyUI Settings Menu Reference
|
||||
|
||||
Configure options directly under ComfyUI Settings (⚙ gear icon):
|
||||
|
||||
- **🍃 FlowControl: Enable Civitai Auto-Scraping**: Automatically download model preview images from Civitai via SHA256 file hashes.
|
||||
- **🍃 FlowControl: Civitai API Key**: Optional key for Civitai SHA256 model preview search (`Authorization: Bearer <key>`).
|
||||
- **🍃 FlowControl: Enable TMDB Auto-Scraping**: Automatically download celebrity preview images from TMDB.
|
||||
- **🍃 FlowControl: TMDB API Key**: Optional key for celebrity preview search.
|
||||
- **🍃 FlowControl: Default Pause Queue State on Launch**: Sets whether the queue should start `Unpaused` or `Paused` on boot.
|
||||
- **🍃 FlowControl: Default Pause Queue Mode on Launch**: Sets the default pause behavior (`Finish Active Prompt` vs `Instant Resume Node`).
|
||||
- **🍃 FlowControl: Persistent Queue Restored Launch State**: Controls what happens to crash-recovered queue items (`Start Paused` or `Start Unpaused`).
|
||||
- **🍃 FlowControl: Pause Button Unpaused Color**: Customizable hex color for the toolbar unpaused state.
|
||||
- **🍃 FlowControl: Pause Button Paused Color**: Customizable hex color for the toolbar paused/pausing state.
|
||||
- **🍃 FlowControl: Enable Pause Queue Toolbar Button**: Toggle top action bar Pause & Continue button ON/OFF.
|
||||
- **🍃 FlowControl: Enable Persistent Queue (Auto-Recovery)**: Toggle the real-time crash recovery queue saving on or off.
|
||||
<details>
|
||||
<summary><b>🖼️ Visual Loaders (Civitai & TMDB Duo)</b></summary>
|
||||
|
||||
- **`Enable Custom SaturnNodes Colors`** (`boolean`, *Default: false*): Applies custom amber colors to SaturnNodes on the canvas. When disabled (default), nodes use default ComfyUI colors.
|
||||
- **`Custom Node Title Color`** (`color`, *Default: `#6a3f20`*): Header / title bar color picker for SaturnNodes when custom colors are enabled.
|
||||
- **`Custom Node Body Color`** (`color`, *Default: `#472a15`*): Body background color picker for SaturnNodes when custom colors are enabled.
|
||||
- **`Civitai API Key`** (`text`): Optional key for Civitai SHA256 model preview search (`Authorization: Bearer <key>`).
|
||||
- **`Enable Civitai Auto-Scraping`** (`boolean`, *Default: false*): Automatically download model preview images from Civitai via SHA256 file hashes. (Note: Local SHA256 hash searching always works).
|
||||
- **`TMDB Access Token`** (`text`): Optional key or v4 Read Access Token (`eyJ...`) for actor / media preview search.
|
||||
- **`Enable TMDB Auto-Scraping`** (`boolean`, *Default: false*): Automatically download media preview images from TMDB.
|
||||
- **`Enable LoRA Usage Tracking`** (`boolean`, *Default: true*): Toggle LoRA usage counting and visual rank badges (🔥, Gold, Silver, Bronze) in the picker.
|
||||
- **`Reset Failed Scrapes Cache`** (*Button: `🗑️ Clear Scrapes Cache`*): Clears failed scrape history so Civitai/TMDB can retry downloading missing preview images.
|
||||
- **`Disable Thumbnail Downscaling (Serve Raw)`** (`boolean`, *Default: false*): Serves full-resolution original preview images instead of generating downscaled WebP thumbnails (downscaled WebP thumbnails render faster in the picker and are recommended).
|
||||
- **`Image Loader Metadata Scan Limit`** (`number`, *Default: 128*, 0–2000): Number of newest images in a folder scanned for embedded prompts by the Visual Image Loader (`0` disables extraction).
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>👁️ Live Preview (Floating Latent Dock)</b></summary>
|
||||
|
||||
- **`Enable Floating Latent Dock`** (`boolean`, *Default: true*): Adds the floating picture-in-picture Live Latent preview dock (bottom-right corner) plus a canvas toolbar toggle button.
|
||||
- **`Floating Dock Max Width`** (`slider`, *Default: 280*, 1px steps): Maximum dock width box in pixels (100–2000). Drag the dock's top-left corner to set a free-form width×height box; the dragged dimensions are saved back into this setting, and the dock then snaps to the exact latent aspect ratio without distortion.
|
||||
- **`Floating Dock Max Height`** (`slider`, *Default: 280*, 1px steps): Maximum dock height box in pixels (100–2000). Drag the dock's top-left corner to set a free-form width×height box; the dragged dimensions are saved back into this setting, and the dock then snaps to the exact latent aspect ratio without distortion.
|
||||
- **`Auto-Hide Dock When Idle`** (`boolean`, *Default: false*): Hides the dock while idle and reveals it automatically when sampling begins.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🔄 Prompt Iterator</b></summary>
|
||||
|
||||
- **`Clear State on Launch`** (`boolean`, *Default: false*): Privacy toggle to empty `prompt_iterator_state.json` on ComfyUI startup.
|
||||
- **`Reset All Queues`** (*Button: `🔄 Reset All Queues`*): Immediately empties all active prompt queues and resets iterator state.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>📋 Prompt Actions</b></summary>
|
||||
|
||||
- **`Show "Copy Prompt" Button on Images`** (`boolean`, *Default: true*): Shows the 📋 "Copy Prompt" overlay action button when hovering over generated images in the Assets / History pane and preview nodes.
|
||||
- **`Show Right-Click "Copy Prompt" Menu Action`** (`boolean`, *Default: true*): Adds "📋 Copy Prompt" to node right-click context menus.
|
||||
- **`Show "Inspect Asset" (Zoom) Button on Images`** (`boolean`, *Default: false*): Restores the 🔍 "Inspect asset" (zoom in) button directly onto image cards in the Assets pane next to Download and Copy Prompt (restoring one-click access moved behind the 3-dots menu in newer ComfyUI versions).
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>⏸️ Pause Controls</b></summary>
|
||||
|
||||
- **`Default State on Launch`** (`combo`, *Default: `Running`*): Sets whether the queue starts `Paused` or `Running` on boot.
|
||||
- **`Default Pause Action`** (`combo`, *Default: `Finish Active Prompt`*): Sets default pause behavior (`Finish Active Prompt` vs `Instant Resume Node`).
|
||||
- **`Persist Pause State Across Restarts`** (`boolean`, *Default: false*): Restores the last manual Pause/Continue state and pause mode on the next start, overriding the launch defaults.
|
||||
- **`Enable Top Toolbar Button`** (`boolean`, *Default: true*): Toggle top action bar Pause & Continue button group ON/OFF.
|
||||
- **`Toolbar Button Unpaused Color`** (`text`, *Default: `#16a34a`*): Hex color for the toolbar unpaused/running state.
|
||||
- **`Toolbar Button Paused Color`** (`text`, *Default: `#ea580c`*): Hex color for the toolbar paused state.
|
||||
- **`Enable System Tray Icon`** (`boolean`, *Default: false*): Displays an OS system tray icon with real-time queue status colors and outside-browser controls.
|
||||
- **`Allow Process Management (Restart / Shutdown)`** (`boolean`, *Default: false*): Opt-in authorization allowing server restart and shutdown actions from the SaturnNodes power controls. Disabled by default for maximum security.
|
||||
- **`Visible Queue Items (Height)`** (`slider`, *Default: automatic per-monitor*, 1px steps, 0–1200): Limits the visible height of the queue overlay panel in pixels (`0` = native). The overlay's bottom edge can also be dragged directly on screen; sizes are remembered per monitor.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>💾 Persistent Queue</b></summary>
|
||||
|
||||
- **`Enable Persistent Queue (Auto-Recovery)`** (`boolean`, *Default: true*): Automatically saves unfinished queue items and restores them after restart/crash.
|
||||
- **`Recovery Launch State`** (`combo`, *Default: `Match Default`*): Override launch state when restored queue items are recovered on startup (`Match Default`, `Force Paused`, `Force Running`).
|
||||
- **`Max Restore Retries`** (`number`, *Default: 0*, 0–10): How many additional launches may retry a queue item that failed to restore before it is dropped. `0` drops failures immediately.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🖼️ Assets & History Restore</b></summary>
|
||||
|
||||
- **`Restore Assets on Launch`** (`boolean`, *Default: true*): Automatically restores your latest generated images into the Assets / History pane on startup.
|
||||
- **`Restored Assets Count`** (`number`, *Default: 64*): The number of newest images from the output folder to populate into the Assets pane.
|
||||
- **`Scan Output Subfolders Recursively`** (`boolean`, *Default: true*): Include images from nested output subfolders, not just the top-level output folder.
|
||||
- **`Show Generation Time on Restored Assets`** (`boolean`, *Default: true*): Uses SaturnNodes' recorded generation timings (from `generation_times.json`) for real execution times on restored history entries. Images generated before this version have no recorded time: their file mtime is still used for chronology, but no duration is shown.
|
||||
- **`Restore Now`** (*Button: `🖼️ Restore Now`*): Immediately rescan the output folder and add the newest assets without restarting ComfyUI.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🩺 Diagnostics & Debug</b></summary>
|
||||
|
||||
- **`Export Debug Profile`** (*Button: `📥 Export Debug Profile`*): Exports non-sensitive system environment details (OS, Python, PyTorch, SaturnNodes settings, local cache counts) to a JSON file to share when reporting bugs or requesting assistance. Sensitive API keys and tokens are never exported.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🎨 Batch Queue Visuals</b></summary>
|
||||
|
||||
- **`Show Batch Queue 1D Lines`** (`boolean`, *Default: true*): Renders a 1D git-graph style colored line on the left side of queued items in the queue list.
|
||||
- **36 Cycling Colors:** Every queued batch is assigned a vibrant, high-contrast color from a 36-color sequence.
|
||||
- **Clean 1D Line:** The first item starts with a top curve (`┌`), middle items are straight lines (`│`), and the last item ends with a bottom curve (`└`). Standalone 1-item batches are rounded at both ends (`(`).
|
||||
- **In-between Infiltration Detection:** If a different prompt is queued in-between items of a batch, the interrupted batch maintains its open line ends without false curves, while the inserted item is highlighted with its own bracket.
|
||||
- **Works Standalone or with PersistentQueue:** Operates client-side via `localStorage` when standalone, and syncs with `PersistentQueue` to retain batch relationships after server restarts.
|
||||
- **`Batch Graph Snapshot Guard`** (`boolean`, *Default: true*): Snapshots prompt text, LoRA selections and node inputs when you click Queue (including batches queued behind an in-flight submission), ensuring later canvas edits do not mutate queued prompts. Random seeds continue to randomize.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🛡️ Security & Script Execution</b></summary>
|
||||
|
||||
- **`Allow Local File Execution`** (`boolean`, *Default: false*): Master toggle controlling whether the *Run Local File* node is permitted to execute local scripts or executables. Disabled by default for operator safety.
|
||||
- **Strict Confinement**: Execution is strictly limited to the `ComfyUI/scripts/` directory. Absolute paths and directory traversal (`..`) are permanently blocked.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>🧩 Node Management</b></summary>
|
||||
|
||||
- One enable/disable toggle per SaturnNodes node (e.g. `Loaders / LoRA Loader (Folder)`, `Utils / Prompt Counter`): hides individual nodes from the node search menu. Disabled nodes remain registered, so existing workflows keep loading unchanged — they are only hidden from the Add Node menu (requires a server restart to take effect).
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
## 🔒 Security & Privacy
|
||||
|
||||
- All API keys and settings are stored locally in `.env` (excluded from git commits via `.gitignore`).
|
||||
- No hardcoded personal paths, keys, or credentials.
|
||||
- Atomic file writes (`.tmp` + `os.replace`) prevent JSON corruption during sudden crashes.
|
||||
For the comprehensive engineering standard, threat model, and review precedents enforced by the ComfyUI-Manager maintainers, see the full **[ComfyUI Custom Node Security & Compliance Rulebook](SECURITY_AUDIT_PLAYBOOK.md)**.
|
||||
|
||||
- **Local Runner Sandbox & Operator Authorization Gate:** The `Run Local File` node enforces single-use interactive operator authorization (`⚡ Authorize Run (5 min)`), strictly confines execution to `ComfyUI/scripts/`, never uses `shell=True`, terminates whole process trees on timeout or cancellation, and **never persists authorization inside workflow JSON or image metadata**.
|
||||
- **CSRF Defense & Fetch Metadata Validation:** All administrative endpoints (`/saturnnodes/settings`, `/saturnnodes/power/*`, `/saturnnodes/local_runner/*`, `/saturnnodes/scrapes/clear`, `/saturnnodes/pause_queue/*`, `/saturnnodes/persistent_queue/*`, `/saturnnodes/image_loader/*`, `/saturnnodes/folder_lora_loader/*`) enforce loopback checks, `Sec-Fetch-Site` header checks (blocking cross-site requests from visited web pages), and cryptographic per-session `X-SaturnNodes-CSRF-Token` headers.
|
||||
- **Two-Step Ephemeral Ticket Handshake for Power Actions:** Process restart and shutdown controls are disabled by default (`ALLOW_PROCESS_MANAGEMENT=false`), require CSRF authentication, and execute only after a single-use 30-second ephemeral ticket confirmation handshake.
|
||||
- **Universal Loopback Protection:** All administrative and queue endpoints strictly verify loopback origin (`127.0.0.1` / `::1`), rejecting external network requests with `403 Forbidden`.
|
||||
- **Directory Traversal Protection:** Image loading and thumbnail routes are strictly confined to ComfyUI `input`, `output`, and `temp` directories with traversal guards.
|
||||
- **Optional Headless Dependencies:** System tray desktop integration (`pystray`) is an optional dependency, ensuring cloud containers (Docker, Colab, RunPod) run with zero dependency errors.
|
||||
- **Opt-In Scraping:** Civitai and TMDB network scraping are opt-in and disabled by default.
|
||||
- **Zero Runtime Package Installs:** No `install.py` or dynamic `pip` execution; all standard dependencies are cleanly declared in `requirements.txt`.
|
||||
- **Local Environment:** All API keys and settings are stored locally in `.env` inside `user/default/SaturnNodes/`.
|
||||
- **Atomic File Writes:** Safe atomic writes (`.tmp` + `os.replace`) prevent JSON corruption during sudden crashes.
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Installation & Setup
|
||||
|
||||
**Option A — ComfyUI Manager / Registry (recommended):**
|
||||
|
||||
1. Open ComfyUI Manager → **Install Custom Nodes** and search for `SaturnNodes`, or run:
|
||||
```bash
|
||||
comfy node install saturnnodes
|
||||
```
|
||||
2. Restart ComfyUI.
|
||||
|
||||
**Option B — Manual (git clone):**
|
||||
|
||||
1. Open your terminal and navigate to your ComfyUI `custom_nodes` directory:
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/KOFiblto/ComfyUI-FlowControl.git
|
||||
git clone https://github.com/KOFiblto/ComfyUI-SaturnNodes.git
|
||||
```
|
||||
2. Restart ComfyUI.
|
||||
3. Open ComfyUI Settings (⚙ gear menu) to configure API keys or feature toggles if desired.
|
||||
|
||||
**After install (either option):** open ComfyUI Settings (⚙ gear menu) to configure API keys or feature toggles if desired, and see the [Complete Walkthrough](WALKTHROUGH.md) for usage recipes.
|
||||
|
||||
---
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! Please check out [CONTRIBUTING.md](CONTRIBUTING.md) for local development guidelines, security policy, and coding standards.
|
||||
|
||||
## 📜 License
|
||||
|
||||
Distributed under the [MIT License](LICENSE).
|
||||
|
||||
+27
-4
@@ -2,7 +2,7 @@
|
||||
|
||||
## Reporting Security Issues
|
||||
|
||||
We take the security of ComfyUI-FlowControl seriously. If you discover a security vulnerability or sensitive key exposure, please **DO NOT** open a public issue.
|
||||
We take the security of ComfyUI-SaturnNodes seriously. If you discover a security vulnerability or sensitive key exposure, please **DO NOT** open a public issue.
|
||||
|
||||
Instead, please report security vulnerabilities directly to the maintainer via GitHub private vulnerability reporting.
|
||||
|
||||
@@ -10,9 +10,32 @@ Instead, please report security vulnerabilities directly to the maintainer via G
|
||||
|
||||
| Version | Supported |
|
||||
| ------- | ------------------ |
|
||||
| 2.x | :white_check_mark: |
|
||||
| 1.x | :white_check_mark: |
|
||||
|
||||
## Security Best Practices
|
||||
## Security Architecture & Best Practices
|
||||
|
||||
- **CSRF Defense & Fetch Metadata Validation**: All administrative endpoints under `/saturnnodes/*` enforce multi-layer CSRF protection via `is_authenticated_local_request(request)`.
|
||||
- Enforces loopback origin checks (`127.0.0.1` / `::1`), rejecting remote calls with `403 Forbidden`.
|
||||
- Rejects `cross-site` requests via `Sec-Fetch-Site` header checks, preventing external malicious web pages visited by the operator from triggering actions.
|
||||
- Verifies `Origin` and `Referer` headers against loopback hostnames.
|
||||
- Requires cryptographically random session tokens in the `X-SaturnNodes-CSRF-Token` header (`secrets.compare_digest`), prompting CORS preflight rejection if an external origin attempts unauthorized requests.
|
||||
- **Two-Step Ephemeral Ticket Handshake for Power Actions**:
|
||||
- Server restart and shutdown endpoints require `ALLOW_PROCESS_MANAGEMENT=true` in settings.
|
||||
- Execution requires a two-step handshake: the client requests a single-use ephemeral ticket (`POST /saturnnodes/power/request_token`) valid for 30 seconds.
|
||||
- The confirmation endpoint (`POST /saturnnodes/power/confirm_action`) consumes the single-use ticket in server memory before executing restart or shutdown, preventing arbitrary or cross-site triggers.
|
||||
- **Strict Local Script Confinement**:
|
||||
- `RunLocalFileNode` strictly confines executable scripts and binaries to `ComfyUI/scripts/`.
|
||||
- Absolute paths (e.g. `C:\`, `/etc/`, `\\server`) and directory traversal sequences (`..`) are strictly rejected.
|
||||
- Paths are verified via canonical `os.path.realpath` and `os.path.commonpath` checks against the scripts directory. Unsafe path bypasses have been completely removed.
|
||||
- **In-Memory Operator Authorization Gate**:
|
||||
- Local file execution defaults to disabled (`ENABLE_LOCAL_FILE_EXECUTION = false`).
|
||||
- `security_consent` input widget was removed from `INPUT_TYPES` to ensure authorization parameters NEVER travel within workflow JSON or embedded PNG metadata.
|
||||
- Workflows loaded from the internet or images dropped onto the canvas cannot execute local scripts without explicit manual authorization.
|
||||
- Execution requires an interactive user action on the canvas node via `[⚡ Authorize Run]`, which registers single-use authorization exclusively in server RAM with a 5-minute TTL.
|
||||
- **Filesystem Path Confinement**: Image loading (`VisualImageLoader`) and thumbnail file streaming endpoints are strictly confined to standard ComfyUI `input`, `output`, and `temp` directories with path canonicalization and directory traversal guards.
|
||||
- **Opt-In Network Requests**: External metadata and preview scraping (Civitai / TMDB) are completely opt-in and disabled by default.
|
||||
- **No Dynamic Package Installation**: Runtime `pip` execution (`install.py` / `subprocess`) is eliminated; all dependencies are transparently listed in `requirements.txt`.
|
||||
- **Local Environment**: All API keys (Civitai / TMDB) are stored locally in `.env` and are strictly excluded from git tracking via `.gitignore`. Sensitive keys are masked in server log outputs.
|
||||
- **Atomic File Writes**: Atomic file operations (`.tmp` + `os.replace`) prevent JSON file corruption on sudden system shutdown or crash.
|
||||
|
||||
- All API keys (Civitai / TMDB) are stored locally in `.env` and are strictly excluded from git tracking via `.gitignore`.
|
||||
- API keys are masked in log outputs (e.g. `Bearer civ_*****`).
|
||||
|
||||
@@ -0,0 +1,435 @@
|
||||
# 🛡️ ComfyUI Custom Node Security, Compliance & Architecture Rulebook
|
||||
*The Authoritative Engineering Standard & Audit Playbook Derived from ComfyUI-Manager PR #3200 Review Precedents*
|
||||
|
||||
---
|
||||
|
||||
## 📑 Table of Contents
|
||||
1. [Core Threat Model & The Golden Rule](#1-core-threat-model--the-golden-rule)
|
||||
2. [Case Study: ComfyUI-Manager PR #3200 Review Precedents](#2-case-study-comfyui-manager-pr-3200-review-precedents)
|
||||
3. [The 8 Core Security Rules (With Code Examples)](#3-the-8-core-security-rules)
|
||||
- [Rule 1: Filesystem Confinement & Traversal Prevention](#rule-1-filesystem-confinement--traversal-prevention)
|
||||
- [Rule 2: Destructive File Actions & Deletion Guards](#rule-2-destructive-file-actions--deletion-guards)
|
||||
- [Rule 3: Process Execution & The Workflow Serialization Trap](#rule-3-process-execution--the-workflow-serialization-trap)
|
||||
- [Rule 4: Dynamic Code Evaluation & Deserialization](#rule-4-dynamic-code-evaluation--deserialization)
|
||||
- [Rule 5: Web Endpoints & The Same-Machine Browser Fallacy (CSRF)](#rule-5-web-endpoints--the-same-machine-browser-fallacy-csrf)
|
||||
- [Rule 6: Two-Step Ephemeral Ticket Handshakes for Power Actions](#rule-6-two-step-ephemeral-ticket-handshakes-for-power-actions)
|
||||
- [Rule 7: Outbound Network Requests & SSRF Prevention](#rule-7-outbound-network-requests--ssrf-prevention)
|
||||
- [Rule 8: Configuration Storage, Isolation & Injection Defense](#rule-8-configuration-storage-isolation--injection-defense)
|
||||
4. [Ecosystem & Registry Packaging Compliance](#4-ecosystem--registry-packaging-compliance)
|
||||
- [Rule 9: Zero Runtime Package Installations (`install.py` Prohibition)](#rule-9-zero-runtime-package-installations-installpy-prohibition)
|
||||
- [Rule 10: PyTorch Exclusion from Dependencies](#rule-10-pytorch-exclusion-from-dependencies)
|
||||
- [Rule 11: Headless & Cloud Container Resilience](#rule-11-headless--cloud-container-resilience)
|
||||
- [Rule 12: Core Monkey-Patching Isolation](#rule-12-core-monkey-patching-isolation)
|
||||
5. [Frontend Hygiene & Canvas Safety](#5-frontend-hygiene--canvas-safety)
|
||||
- [Rule 13: Dual Frontend Architecture (V1 LiteGraph & V2 Vue UI)](#rule-13-dual-frontend-architecture-v1-litegraph--v2-vue-ui)
|
||||
- [Rule 14: Canvas Resize Loops (`computeSize` & `node.size`)](#rule-14-canvas-resize-loops-computesize--nodesize)
|
||||
- [Rule 15: Non-Serialized UI Widgets (`{ serialize: false }`)](#rule-15-non-serialized-ui-widgets--serialize-false-)
|
||||
6. [Auditor & Maintainer Quick Conformance Matrix](#6-auditor--maintainer-quick-conformance-matrix)
|
||||
|
||||
---
|
||||
|
||||
## 1. Core Threat Model & The Golden Rule
|
||||
|
||||
### 1.1 The Golden Rule
|
||||
> **"Treat every input, widget values included, as attacker-controlled, because it can come from a shared workflow or a direct `/prompt` call."**
|
||||
> — *ltdrdata (ComfyUI-Manager Lead Maintainer, Comfy-Org)*
|
||||
|
||||
### 1.2 The 4 Attack Vectors Unique to ComfyUI
|
||||
1. **Shared Workflow Vectors (.json & PNG Metadata):**
|
||||
ComfyUI workflows are freely shared across community hubs (Civitai, Reddit, Discord, GitHub) as raw JSON files or embedded inside PNG image generation metadata. Loading a poisoned image instantly sets all node widgets on the canvas.
|
||||
2. **The Same-Machine Browser Exploit (Drive-by Localhost CSRF):**
|
||||
Any website an operator visits in their regular browser can execute background JavaScript (`fetch('http://localhost:8188/...')`). Because this request originates from the operator's local computer, primitive IP loopback checks (`127.0.0.1`) pass automatically unless strict fetch metadata and CSRF headers are validated.
|
||||
3. **Direct Headless API Injection (`POST /prompt`):**
|
||||
Attackers or automated scripts can bypass the graphical interface completely by sending malicious graph specifications directly to the server's `/prompt` endpoint.
|
||||
4. **Host User Privileges:**
|
||||
ComfyUI processes run with the full security privileges of the logged-in operating system user. Any sandbox breakout, unchecked script runner, or arbitrary file deletion results in immediate compromise of personal files or host systems.
|
||||
|
||||
---
|
||||
|
||||
## 2. Case Study: ComfyUI-Manager PR #3200 Review Precedents
|
||||
|
||||
During the official registration of **`ComfyUI-SaturnNodes`** in [Comfy-Org/ComfyUI-Manager#3200](https://github.com/Comfy-Org/ComfyUI-Manager/pull/3200), lead maintainer `@ltdrdata` conducted a three-phase security review. This exchange set definitive precedents for all modern custom nodes:
|
||||
|
||||
### Round 1: Unauthenticated Power Routes & Arbitrary File Reads
|
||||
* **Maintainer Finding:** Exposing network routes that restart/shutdown the server, stream files by absolute path, or write to `.env` without access controls is strictly forbidden.
|
||||
* **Precedent:** File reads must be confined to fixed base directories. Power management and settings routes must be local-only and strictly constrained.
|
||||
|
||||
### Round 2: The Same-Machine Fallacy & The Workflow Serialization Trap
|
||||
* **Maintainer Finding:**
|
||||
1. *Same-Machine Check Failure:* A basic IP loopback check (`request.remote == "127.0.0.1"`) is insufficient. Any open browser tab makes requests from `localhost`, allowing malicious web pages to trigger restarts, shutdowns, or settings changes.
|
||||
2. *Escape Hatches Forbidden:* Settings toggles such as `ALLOW_ANY_SCRIPT_PATH` that bypass folder confinement must be completely removed.
|
||||
3. *In-Workflow Consent Fallacy:* A `security_consent` boolean widget in `INPUT_TYPES` defaults to `False` in code, but a shared malicious workflow arrives with `security_consent: True` already serialized. Execution consent must never travel with the graph.
|
||||
* **Precedent:**
|
||||
- Routes must validate `Sec-Fetch-Site` (blocking `cross-site`) and verify cryptographic CSRF tokens.
|
||||
- Process execution must require interactive, in-memory human authorization on the canvas node via an ephemeral RAM token.
|
||||
|
||||
### Round 3: Folder Deletion Sandbox Bypass
|
||||
* **Maintainer Finding:** `sanitize_folder_path` allowed relative paths (`..`) and absolute paths. In a folder watcher node with an optional `delete_image` flag, an attacker could supply `folder: "C:/Windows"` or `../../` and delete arbitrary user files.
|
||||
* **Precedent:** All path-handling functions must reject absolute paths, reject `..`, and verify canonical containment against approved ComfyUI directories using `os.path.commonpath`.
|
||||
|
||||
---
|
||||
|
||||
## 3. The 8 Core Security Rules
|
||||
|
||||
### Rule 1: Filesystem Confinement & Traversal Prevention
|
||||
|
||||
Nodes must never read, write, or list files outside designated ComfyUI directories (`input/`, `output/`, `temp/`, and `models/`).
|
||||
|
||||
#### ❌ Non-Compliant Patterns
|
||||
```python
|
||||
# VULNERABLE: Accepts arbitrary absolute paths or directory traversal
|
||||
def load_image(self, folder, filename):
|
||||
filepath = os.path.join(folder, filename)
|
||||
with open(filepath, "rb") as f: # Attacker loads C:\Users\user\.ssh\id_rsa
|
||||
return f.read()
|
||||
```
|
||||
|
||||
#### ✅ Compliant Pattern
|
||||
```python
|
||||
import os
|
||||
import re
|
||||
|
||||
def sanitize_folder_path(folder_input: str, allowed_base_dir: str) -> str:
|
||||
"""
|
||||
Strictly confines a folder to an allowed base directory.
|
||||
Rejects absolute paths, drive letters, UNC shares, and parent traversal.
|
||||
"""
|
||||
if not folder_input or not folder_input.strip():
|
||||
return allowed_base_dir
|
||||
|
||||
raw_path = folder_input.strip()
|
||||
|
||||
# Reject Windows drive letters (C:) and UNC network paths (// or \\)
|
||||
if re.match(r'^[A-Za-z]:', raw_path) or raw_path.startswith(("\\\\", "//")):
|
||||
raise ValueError(f"Absolute or network paths are forbidden: {raw_path}")
|
||||
|
||||
# Reject root slashes and directory traversal
|
||||
if raw_path.startswith(("/", "\\")) or ".." in raw_path.replace("\\", "/").split("/"):
|
||||
raise ValueError(f"Directory traversal sequences are forbidden: {raw_path}")
|
||||
|
||||
# Canonicalize and verify commonpath boundary
|
||||
canonical_base = os.path.realpath(allowed_base_dir)
|
||||
target_path = os.path.realpath(os.path.join(canonical_base, raw_path))
|
||||
|
||||
if os.path.commonpath([canonical_base, target_path]) != canonical_base:
|
||||
raise ValueError(f"Path escapes allowed sandbox: {target_path}")
|
||||
|
||||
return target_path
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Rule 2: Destructive File Actions & Deletion Guards
|
||||
|
||||
Any node that deletes or overwrites files (`os.remove`, `os.unlink`, `shutil.rmtree`) must verify containment immediately prior to invocation.
|
||||
|
||||
#### ❌ Non-Compliant Patterns
|
||||
```python
|
||||
# VULNERABLE: Trusting earlier sanitization or allowing unverified deletion
|
||||
def cleanup_cache(self, user_supplied_path):
|
||||
if os.path.exists(user_supplied_path):
|
||||
os.remove(user_supplied_path) # Arbitrary file deletion
|
||||
```
|
||||
|
||||
#### ✅ Compliant Pattern
|
||||
```python
|
||||
def safe_delete_image(target_file_path: str, allowed_folders: list[str]) -> bool:
|
||||
"""
|
||||
Verifies target file is strictly inside input/output before deletion.
|
||||
"""
|
||||
real_target = os.path.realpath(target_file_path)
|
||||
if not os.path.isfile(real_target):
|
||||
return False
|
||||
|
||||
is_confined = any(
|
||||
os.path.commonpath([os.path.realpath(base), real_target]) == os.path.realpath(base)
|
||||
for base in allowed_folders
|
||||
)
|
||||
|
||||
if not is_confined:
|
||||
raise PermissionError(f"Security Violation: Refusing to delete file outside sandbox: {real_target}")
|
||||
|
||||
os.remove(real_target)
|
||||
return True
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Rule 3: Process Execution & The Workflow Serialization Trap
|
||||
|
||||
Executing external binaries or scripts is high risk. You must adhere to the following principles:
|
||||
|
||||
1. **The Workflow Serialization Trap:** Never include an authorization checkbox (e.g. `"security_consent": ("BOOLEAN", {"default": False})`) in `INPUT_TYPES`. Attackers distribute workflows with `"security_consent": True` already serialized.
|
||||
2. **Ephemeral RAM Authorization:** Authorization must occur via an in-memory session token generated when the user physically clicks a canvas UI button (e.g. `[⚡ Authorize Run]`).
|
||||
3. **Execution Confinement:** Scripts must reside strictly within a dedicated directory (e.g., `ComfyUI/scripts/`).
|
||||
4. **Never Use `shell=True`:** Always tokenize arguments (`shlex.split` or list) and execute with `shell=False`.
|
||||
|
||||
#### ❌ Non-Compliant Patterns
|
||||
```python
|
||||
# VULNERABLE: shell=True, arbitrary paths, and consent stored in workflow JSON
|
||||
class VulnerableRunner:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"command": ("STRING", {"default": "python script.py"}),
|
||||
"user_consent": ("BOOLEAN", {"default": False}), # Poisoned via JSON!
|
||||
}
|
||||
}
|
||||
def run(self, command, user_consent):
|
||||
if user_consent:
|
||||
subprocess.Popen(command, shell=True) # Remote Code Execution
|
||||
```
|
||||
|
||||
#### ✅ Compliant Pattern
|
||||
```python
|
||||
import subprocess
|
||||
import shlex
|
||||
import time
|
||||
import os
|
||||
|
||||
# Stored strictly in volatile Python memory (RAM), never serialized to disk
|
||||
ACTIVE_RUN_AUTHORIZATIONS: dict[str, float] = {}
|
||||
|
||||
def authorize_node_execution(node_id: str, ttl_seconds: int = 300) -> str:
|
||||
"""Invoked only by interactive UI button via authenticated CSRF route."""
|
||||
token = os.urandom(16).hex()
|
||||
ACTIVE_RUN_AUTHORIZATIONS[node_id] = time.time() + ttl_seconds
|
||||
return token
|
||||
|
||||
def execute_confined_script(script_name: str, args_list: list[str], node_id: str, scripts_dir: str):
|
||||
# 1. Verify in-memory authorization
|
||||
expiry = ACTIVE_RUN_AUTHORIZATIONS.pop(node_id, 0)
|
||||
if time.time() > expiry:
|
||||
raise PermissionError("Execution denied: Interactive operator authorization required on canvas.")
|
||||
|
||||
# 2. Confine target executable to scripts/ directory
|
||||
real_scripts_base = os.path.realpath(scripts_dir)
|
||||
target_script = os.path.realpath(os.path.join(real_scripts_base, script_name))
|
||||
|
||||
if os.path.commonpath([real_scripts_base, target_script]) != real_scripts_base:
|
||||
raise ValueError("Script escape detected.")
|
||||
|
||||
if not os.path.isfile(target_script):
|
||||
raise FileNotFoundError(f"Script not found: {script_name}")
|
||||
|
||||
# 3. Safe execution with shell=False
|
||||
cmd = ["python", target_script] + args_list
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
shell=False,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True
|
||||
)
|
||||
return proc.communicate(timeout=60)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Rule 4: Dynamic Code Evaluation & Deserialization
|
||||
|
||||
Never pass widget string inputs into `eval()`, `exec()`, or dynamic `importlib` calls.
|
||||
|
||||
#### ❌ Non-Compliant Patterns
|
||||
```python
|
||||
# VULNERABLE: Direct eval allows __import__('os').system('...')
|
||||
def calculate_aspect(self, formula_string):
|
||||
return eval(formula_string)
|
||||
```
|
||||
|
||||
#### ✅ Compliant Pattern
|
||||
```python
|
||||
import ast
|
||||
import operator
|
||||
|
||||
SAFE_OPERATORS = {
|
||||
ast.Add: operator.add,
|
||||
ast.Sub: operator.sub,
|
||||
ast.Mult: operator.mul,
|
||||
ast.Div: operator.truediv,
|
||||
}
|
||||
|
||||
def safe_eval_math(expression: str) -> float:
|
||||
"""Parses mathematical expression using an Abstract Syntax Tree (AST)."""
|
||||
node = ast.parse(expression, mode='eval').body
|
||||
def _eval(node):
|
||||
if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
|
||||
return node.value
|
||||
if isinstance(node, ast.BinOp) and type(node.op) in SAFE_OPERATORS:
|
||||
return SAFE_OPERATORS[type(node.op)](_eval(node.left), _eval(node.right))
|
||||
raise ValueError(f"Unsafe expression: {expression}")
|
||||
return float(_eval(node))
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Rule 5: Web Endpoints & The Same-Machine Browser Fallacy (CSRF)
|
||||
|
||||
Because any malicious web page the operator visits in another tab can target `http://127.0.0.1:8188`, endpoints that mutate state or trigger actions must defend against cross-site requests.
|
||||
|
||||
#### Required Multi-Layer Defense:
|
||||
1. **Loopback IP Verification:** Peer address must be `127.0.0.1`, `::1`, or `localhost`.
|
||||
2. **Fetch Metadata Inspection:** `Sec-Fetch-Site` header must **not** be `cross-site`.
|
||||
3. **Cryptographic CSRF Header:** Require a per-session secret token in a custom header (e.g. `X-SaturnNodes-CSRF-Token`). Browsers forbid cross-origin pages from sending custom headers without CORS preflight approval.
|
||||
|
||||
#### ✅ Compliant Pattern
|
||||
```python
|
||||
import secrets
|
||||
from aiohttp import web
|
||||
|
||||
SESSION_CSRF_TOKEN = secrets.token_hex(32)
|
||||
|
||||
def is_authenticated_local_request(request: web.Request) -> bool:
|
||||
"""
|
||||
Blocks remote network callers AND malicious cross-site browser requests.
|
||||
"""
|
||||
# 1. IP Loopback Check
|
||||
peername = request.transport.get_extra_info('peername')
|
||||
client_ip = peername[0] if peername else ""
|
||||
if client_ip not in ("127.0.0.1", "::1", "localhost"):
|
||||
return False
|
||||
|
||||
# 2. Browser Fetch Metadata Defense (blocks drive-by background fetches)
|
||||
sec_fetch_site = request.headers.get("Sec-Fetch-Site", "")
|
||||
if sec_fetch_site in ("cross-site",):
|
||||
return False
|
||||
|
||||
# 3. Custom Header & Secret Verification (triggers CORS preflight rejection)
|
||||
token = request.headers.get("X-SaturnNodes-CSRF-Token", "")
|
||||
if not token or not secrets.compare_digest(token, SESSION_CSRF_TOKEN):
|
||||
return False
|
||||
|
||||
return True
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Rule 6: Two-Step Ephemeral Ticket Handshakes for Power Actions
|
||||
|
||||
High-impact administrative actions (restarting ComfyUI, shutting down the host, mass cache purges) must **never** trigger from a single `POST` endpoint.
|
||||
|
||||
#### Handshake Protocol:
|
||||
1. **Disabled by Default:** Master setting `ALLOW_PROCESS_MANAGEMENT = False`.
|
||||
2. **Step 1 (`POST /power/request_token`):** Server validates local CSRF session, generates a single-use token in RAM with a 30-second TTL, and returns it to the client.
|
||||
3. **Step 2 (`POST /power/confirm_action`):** Client provides the single-use token. Server validates, burns the token, and schedules the action.
|
||||
|
||||
---
|
||||
|
||||
### Rule 7: Outbound Network Requests & SSRF Prevention
|
||||
|
||||
If your custom node scrapes images, models, or metadata from the internet (Civitai, HuggingFace, TMDB):
|
||||
|
||||
1. **Opt-In Only:** Background scraping must be disabled by default in settings.
|
||||
2. **Protocol Confinement:** Allow only `http://` and `https://`. Reject `file://`, `ftp://`, and `gopher://`.
|
||||
3. **Private IP & Cloud Metadata Blocking (SSRF):** Reject requests resolving to:
|
||||
- Loopback: `127.0.0.0/8`, `::1`
|
||||
- Private LAN: `10.0.0.0/8`, `172.16.0.0/12`, `192.168.0.0/16`
|
||||
- Link-Local / Cloud Metadata: `169.254.169.254`
|
||||
|
||||
#### ✅ Compliant Pattern
|
||||
```python
|
||||
import ipaddress
|
||||
import socket
|
||||
from urllib.parse import urlparse
|
||||
|
||||
def is_safe_external_url(url: str) -> bool:
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
return False
|
||||
|
||||
hostname = parsed.hostname
|
||||
if not hostname:
|
||||
return False
|
||||
|
||||
# Resolve IP and check for private / loopback / metadata ranges
|
||||
ip_addr = socket.gethostbyname(hostname)
|
||||
ip = ipaddress.ip_address(ip_addr)
|
||||
|
||||
if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved:
|
||||
return False
|
||||
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Rule 8: Configuration Storage, Isolation & Injection Defense
|
||||
|
||||
All user runtime state, custom settings, and caches must be stored cleanly inside ComfyUI's user directory:
|
||||
`ComfyUI/user/default/<YourCustomNodeName>/`
|
||||
|
||||
- **Newline Injection Defense:** When writing to `.env` or config files, strip `\r` and `\n` from both keys and values to prevent setting forgery.
|
||||
- **Atomic Disk Writes:** Write data to a temporary file (`.tmp`) first, then commit via `os.replace` to prevent corrupted files on unexpected shutdowns.
|
||||
- **Log Masking:** Strip or mask API keys (`sk-***`, token hashes) in console output.
|
||||
|
||||
---
|
||||
|
||||
## 4. Ecosystem & Registry Packaging Compliance
|
||||
|
||||
To pass automated and manual review in the official ComfyUI Registry and ComfyUI-Manager:
|
||||
|
||||
### Rule 9: Zero Runtime Package Installations (`install.py` Prohibition)
|
||||
- **Prohibited:** Creating an `install.py` script that invokes `subprocess.run(["pip", "install", ...])` during ComfyUI startup.
|
||||
- **Mandated:** Declare all dependencies cleanly in `requirements.txt` and `pyproject.toml`.
|
||||
|
||||
### Rule 10: PyTorch Exclusion from Dependencies
|
||||
- **Prohibited:** Listing `torch` or `torchvision` in your custom node's `requirements.txt` or `pyproject.toml`.
|
||||
- **Reason:** ComfyUI manages its own PyTorch/CUDA environment. Specifying `torch` causes pip to downgrade or corrupt the user's CUDA acceleration binaries.
|
||||
|
||||
### Rule 11: Headless & Cloud Container Resilience
|
||||
- If using desktop GUI dependencies (e.g. `pystray` system tray icons, Tkinter windows), wrap imports in `try/except ImportError` and provide silent headless fallbacks so cloud instances (RunPod, Colab, Docker) do not crash on launch.
|
||||
|
||||
### Rule 12: Core Monkey-Patching Isolation
|
||||
- If overriding core ComfyUI objects (such as `PromptQueue`), the monkey-patch must be conditionally executed **only** if the corresponding feature flag is enabled in settings. If the user disables the feature, core classes must remain untouched.
|
||||
|
||||
---
|
||||
|
||||
## 5. Frontend Hygiene & Canvas Safety
|
||||
|
||||
### Rule 13: Dual Frontend Architecture (V1 LiteGraph & V2 Vue UI)
|
||||
- All custom UI extensions must verify compatibility with both the classic LiteGraph canvas and the modern ComfyUI Vue V2 interface.
|
||||
- In hooks (`nodeCreated`, `beforeRegisterNodeDef`), verify both `node.comfyClass` and `node.type`.
|
||||
|
||||
### Rule 14: Canvas Resize Loops (`computeSize` & `node.size`)
|
||||
- Never compute widget heights dynamically based on `this.size[1]` inside `computeSize()`. In LiteGraph, this creates an infinite resize loop where the node expands downwards continuously on every draw frame.
|
||||
|
||||
### Rule 15: Non-Serialized UI Widgets (`{ serialize: false }`)
|
||||
- Dynamic preview widgets, action buttons, ranking badges, and status labels must be configured with:
|
||||
```javascript
|
||||
widget.serialize = false;
|
||||
// or widget.options = { serialize: false };
|
||||
```
|
||||
This prevents transient UI states from cluttering the user's saved workflow JSON.
|
||||
|
||||
---
|
||||
|
||||
## 6. Auditor & Maintainer Quick Conformance Matrix
|
||||
|
||||
Use this matrix when auditing your own or third-party ComfyUI custom nodes:
|
||||
|
||||
| Domain | Mandatory Security Requirement | Verification Method |
|
||||
| :--- | :--- | :--- |
|
||||
| **Filesystem** | Absolute paths (`C:`, `/`, `\\`) rejected | Test with `C:\Windows\System32` & `/etc/passwd` |
|
||||
| **Filesystem** | Directory traversal (`..`) rejected | Test with `../../` |
|
||||
| **Filesystem** | Canonical path confinement via `commonpath` | Verify `commonpath([base, target]) == base` |
|
||||
| **Deletion** | Deletions strictly confined to `input/` and `output/` | Verify boundary check immediately before `os.remove` |
|
||||
| **Execution** | `shell=False` enforced on all subprocesses | Inspect all `subprocess.Popen` / `subprocess.run` calls |
|
||||
| **Execution** | Scripts confined strictly to `ComfyUI/scripts/` | Verify realpath check against scripts directory |
|
||||
| **Execution** | Execution consent **absent** from `INPUT_TYPES` | Verify no authorization widgets in workflow JSON |
|
||||
| **Execution** | In-memory interactive operator authorization | Confirm single-use RAM token generated via UI button |
|
||||
| **Network** | Loopback checks on administrative routes | Test request from non-local IP (expects 403) |
|
||||
| **Network** | CSRF defense: `Sec-Fetch-Site` + custom token header | Test request with `Sec-Fetch-Site: cross-site` |
|
||||
| **Network** | Two-step ephemeral ticket for restart/shutdown | Verify token generation + confirmation handshake |
|
||||
| **Network** | SSRF: Private IPs & cloud metadata blocked | Test with `127.0.0.1`, `192.168.1.1`, `169.254.169.254` |
|
||||
| **Network** | Scraping features opt-in (disabled by default) | Verify settings default to `false` |
|
||||
| **Packaging** | No `install.py` / dynamic pip execution | Verify absence of `install.py` |
|
||||
| **Packaging** | `torch` excluded from `requirements.txt` | Verify dependencies list |
|
||||
| **Data** | Data stored in `user/default/<pack>/` | Inspect state and cache file paths |
|
||||
| **Frontend** | UI widgets marked `{ serialize: false }` | Inspect exported workflow JSON for UI clutter |
|
||||
|
||||
---
|
||||
*Maintained under the ComfyUI-SaturnNodes security framework. Conforms strictly to Comfy-Org registry audit policies.*
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
# 🪐 ComfyUI-SaturnNodes Complete Walkthrough
|
||||
|
||||
Welcome to the full guide for ComfyUI-SaturnNodes! This custom node suite is designed to give you unprecedented control over your generations, visual asset management, and queue resilience.
|
||||
|
||||
This walkthrough covers everything from basic setup to advanced workflows across both the new ComfyUI V2 (Vue) and the classic ComfyUI (LiteGraph).
|
||||
|
||||
---
|
||||
|
||||
## 1. Global Features
|
||||
|
||||
### The Pause Queue Toolbar
|
||||
Upon installing the extension, you will notice a new **Green/Orange Pause Button** nestled next to your standard "Queue Prompt" button.
|
||||
- **Green (Running)**: ComfyUI operates normally.
|
||||
- **Orange (Pausing/Paused)**: ComfyUI is holding the queue.
|
||||
|
||||
You have two pause modes, accessible by clicking the small dropdown arrow next to the Pause button:
|
||||
1. **Pause (Finish)**: This will let your currently generating workflow finish completely, but it will intercept and hold any subsequent workflows in the queue until you click Continue.
|
||||
2. **Pause (Instant)**: This tells ComfyUI to freeze *immediately* before executing the next node in your active workflow. You can resume precisely from where you left off later without losing progress!
|
||||
|
||||
### The System Tray Icon (Outside-Browser Queue Control)
|
||||
Want to check your queue status or pause/resume generations while working in other apps (Photoshop, Blender, games) without opening ComfyUI in your browser?
|
||||
1. Open ComfyUI Settings (⚙️ gear icon) and enable **`Enable System Tray Icon`** (search for `SaturnNodes` in the settings search box, or set `ENABLE_TRAY_ICON=true` in `ComfyUI/user/default/SaturnNodes/.env`).
|
||||
2. SaturnNodes will spawn an OS system tray icon in your Windows/macOS/Linux taskbar notification area:
|
||||
- **Green Icon (❚❚)**: Indicates ComfyUI queue is actively running.
|
||||
- **Orange Icon (▶)**: Indicates ComfyUI queue is paused and waiting for you to continue.
|
||||
- **Orange Icon (❚❚)**: Indicates ComfyUI is in the middle of a generation and will pause as soon as the active workflow/node finishes.
|
||||
3. Left-click the tray icon to quickly toggle between **Pause** and **Continue**.
|
||||
4. Right-click the tray icon to choose explicit actions: **Pause (Finish)**, **Pause (Instant)**, **Continue / Resume**, switch default pause modes, or quickly **Open ComfyUI in Browser**.
|
||||
|
||||
### The 1D Git-Graph Batch Queue Visualizer
|
||||
SaturnNodes automatically tracks which workflows are submitted together as a batch (whether through "Batch count" or sequential submission) and renders a clean, purely graphical 1D git-graph line directly to the left of your queued items in the queue list:
|
||||
- **36 High-Contrast Cycling Colors**: Every new batch gets a distinct color from an accessible 36-color sequence so adjacent batches are visually distinguishable at a glance.
|
||||
- **Purely Graphical (No Clutter)**: No numbers or text labels.
|
||||
- **Batch Start (`┌`)**: The top of the line curves inward towards the card and runs straight down.
|
||||
- **Batch Middle (`│`)**: Seamless, continuous straight vertical line.
|
||||
- **Batch End (`└`)**: Runs straight from above and curves inward towards the card at the bottom.
|
||||
- **Single Item (`(`)**: A standalone 1-item batch curves inward at both top and bottom like a bracket.
|
||||
- **In-between Infiltration**: If you use "Queue Front" or insert another prompt between items of an active batch, the original batch maintains its open ends (no false start/end curves), while the inserted prompt is enclosed in its own bracket.
|
||||
- **Works Standalone or with Persistent Queue**: Runs locally via `localStorage`, and synchronizes with the Persistent Queue so batch groupings survive server restarts.
|
||||
- **V1 & V2 Frontend Compatible**: Works in the modern V2 queue sidebar (`JobAssetsList`) and classic V1 queue dialog.
|
||||
|
||||
### The Persistent Queue (Crash Recovery)
|
||||
You don't need to interact with this—it runs silently in the background!
|
||||
Every time you queue a prompt, SaturnNodes saves the entire queue state to your hard drive in real time. If your computer crashes, power goes out, or you accidentally close the terminal, ComfyUI will automatically restore your unfinished queue upon next boot. The **Recovery Launch State** setting controls whether recovery uses your normal launch state (`Match Default`), always pauses (`Force Paused`), or always runs (`Force Running`); use `Force Paused` if you want to inspect the restored queue before it starts.
|
||||
|
||||
The queue file is versioned. Older queue files are upgraded automatically, and if you downgrade to an older SaturnNodes version, a queue file written by a newer version is left read-only instead of being overwritten.
|
||||
|
||||
### Automatic Assets Pane & History Restore
|
||||
In standard ComfyUI, restarting your server wipes the in-memory execution history, leaving your sidebar Assets Pane / History tab completely blank on launch.
|
||||
- SaturnNodes automatically scans your `output` folder (read-only) on server startup and pre-populates the history with your latest **64 images** (or custom count).
|
||||
- Scans run without blocking ComfyUI's queue mutex, include nested output subfolders by default, and understand PNG, JPEG, WebP, ComfyUI graphs, and A1111/Forge `parameters` metadata.
|
||||
- All images are loaded in strict chronological order with embedded workflow and prompt graphs intact.
|
||||
- Features like **Copy Prompt** and **Copy Workflow** immediately work with restored images as if they were generated in your current session!
|
||||
- New generations record their real execution duration in `generation_times.json`; restored images show that time when available. Older images still use their file mtime for ordering, but no invented duration is shown when none was recorded.
|
||||
- Use **`Restore Now`** in Settings to rescan immediately, or adjust **`Restored Assets Count`**, **`Scan Output Subfolders Recursively`**, and **`Show Generation Time on Restored Assets`** (search for `SaturnNodes`).
|
||||
|
||||
### Image Overlay Actions & Prompt Bookmarks
|
||||
Hovering over any image in the ComfyUI Assets / History pane or preview nodes provides powerful quick actions directly on the card:
|
||||
- **📋 Copy Prompt**: Instantly copies the positive prompt text/graph to your clipboard.
|
||||
- **🔖 Save to Prompt Bookmarks**: One-click saves prompt text straight into your Prompt Bookmarks list with robust SVG bookmark iconography.
|
||||
- **🔍 Inspect Asset (Zoom)**: When enabled via the **Show "Inspect Asset" (Zoom) Button on Images** setting, restores the direct full-size zoom inspect button right on the card hover pill bar alongside Download and Copy Prompt (bringing back quick one-click access that newer ComfyUI frontend versions moved into the 3-dots submenu).
|
||||
- **Dynamic Button Pill Styling**: The action button bar dynamically calculates borders and corner rounding (`rounded-l-lg`, `rounded-r-lg`) depending on which actions are enabled in settings, maintaining a seamless native appearance.
|
||||
|
||||
---
|
||||
|
||||
## 2. Advanced Nodes
|
||||
|
||||
### 🪐 ⏸️ Saturn Decision
|
||||
Sometimes you run a complex workflow that takes hours (e.g., upscaling), but you want to check a low-res preview first to ensure it's worth the compute time.
|
||||
1. Drop the **`🪐 ⏸️ Saturn Decision`** node into your workflow between the low-res and high-res stages.
|
||||
2. When the execution reaches this node, your workflow will pause, and a native desktop notification (Windows/macOS/Linux) will alert you!
|
||||
3. You will see a popup UI on your ComfyUI canvas giving you three choices:
|
||||
- **Continue**: Proceeds with the workflow normally.
|
||||
- **Cancel**: Outputs a `True` signal. You can route this into a Switch node to bypass the high-res stage entirely.
|
||||
- **Stop Workflow**: Instantly aborts the generation queue, exactly like clicking the red X in ComfyUI.
|
||||
|
||||
### 🪐 📂 Load Image From Folder
|
||||
An absolute powerhouse for automation. Point this node at any folder on your hard drive, and it will automatically ingest incoming images, convert them into tensors, and delete the original files to keep the folder clean.
|
||||
- **wait_if_folder_is_empty**: If set to `True`, the workflow will patiently pause and poll the folder at your `rescan_interval` until an image arrives. Perfect for linking ComfyUI to external tools like Photoshop or Blender auto-exports!
|
||||
- **Regex Filtering**: Only want it to grab images that start with `render_`? Just use standard regex like `^render_.*\.png$`.
|
||||
|
||||
### 🪐 🖼️ Visual LoRA & Image Loaders
|
||||
Tired of guessing what a LoRA does based on its filename?
|
||||
1. Use the **`🪐 🖼️ Visual LoRA Loader`**. It natively extracts preview images if you place an identically named `.png`/`.jpg` next to your `.safetensors` file.
|
||||
2. If you enable the **Civitai Auto-Scraper** in your ComfyUI Settings, the node will calculate the SHA256 hash of your LoRA and automatically download the official preview thumbnail directly from Civitai!
|
||||
3. **Pretty Name String Output**: The node outputs the parsed, human-readable name of the loaded LoRA (e.g. "Addison Rae V1") by default, allowing you to route it into a text prompt or watermark node. You can easily switch this via the `output_name` toggle on the node to output the raw filename without extension instead.
|
||||
4. **Independent Per-Node Zoom**: Use the `tile_size` slider to adjust preview card sizes to your preference. Zoom levels are saved per node directly inside your workflow JSON, so different visual loaders maintain independent tile sizes across browser reloads.
|
||||
5. **Maintenance Reset Buttons**: The control bar offers `↺ Counter` (reset this node's Sequential/Random LoRA cycle + prompt→LoRA assignments), `🧹 Usage` (clear usage badges), and `🗑 Cache` (delete cached thumbnails). Each asks for confirmation. ⚠️ Only use `↺ Counter` when the queue is empty — resetting mid-batch restarts the LoRA cycle from the first selected LoRA and can mismatch prompts against LoRAs.
|
||||
|
||||
### 🪐 ⏱️ Recent Outputs
|
||||
Want to quickly reference your recent generations without digging through Windows Explorer? The **`🪐 ⏱️ Recent Outputs`** node cycles chronologically through your `output/` directory, pulling your freshest creations right back into the workflow canvas.
|
||||
|
||||
### 🪐 📝 Prompt Counter & 🪐 🔄 Prompt Queue Iterator
|
||||
Managing multi-prompt queues and setting the exact ComfyUI batch size has never been easier:
|
||||
1. **`🪐 📝 Prompt Counter`**: Place this multiline text node in your workflow and paste all your prompts.
|
||||
- As you type or paste, the live badge instantly updates (e.g. `24 Prompts`), so you know exactly what batch size to set before clicking Queue!
|
||||
- Under Advanced Options, select your delimiter (`>1 Empty Line`, `Newline`, `>2 Empty Lines`, or `Custom Regex`).
|
||||
- If `Custom Regex` is selected, enter any custom delimiter pattern (such as `\n---\n`). The regex field is dynamically enabled only when Custom Regex is active.
|
||||
2. **`🪐 🔄 Prompt Queue Iterator`**: Connect the `STRING` output of the Prompt Counter into the Prompt Queue Iterator to sequentially pop and execute each prompt run-by-run.
|
||||
- Now supports `Custom Regex` in addition to empty line and newline separators.
|
||||
- Choose **`Text & Node`** tracking to keep independent pointers per prompt text/splitter, or **`Node Only`** to keep one pointer per node while you edit prompts.
|
||||
- `Sequential (Stop on End)` exposes an **`exhausted`** boolean so a workflow can stop after the last prompt.
|
||||
- Includes live progress counter widget (`Run X / Y`), an index reset button (`Reset Counter (0)`, now with a confirmation explaining that it restarts the run pointer and clears crash-recovery assignments/shuffle state — it does **not** empty the queue state file; use Settings → `🔄 Reset All Queues` for a full wipe), and direct state file access (`Open Queue File`).
|
||||
|
||||
---
|
||||
|
||||
## 3. Formatting Folder Paths
|
||||
Several nodes across this suite (e.g. *Load Image From Folder*, *Recent Outputs*, *Visual Image Loader*, *Visual LoRA Loader*) take a `folder` or `folder_path` string as input. Here is how they work:
|
||||
|
||||
- **Empty String**: Leaving a field completely empty (or spaces) will default to the sensible native directory for that node (e.g., leaving it empty on *Recent Outputs* will automatically target ComfyUI's root `output/` directory).
|
||||
- **Safe Directory Containment**: For security, image loading and visual thumbnail nodes are confined to ComfyUI `input`, `output`, and `temp` directories with path traversal protection.
|
||||
- **Subfolder Paths**: You can specify subdirectories like `batch_01` or `renders/portraits` within the allowed ComfyUI directories.
|
||||
- **Wildcard Filters (`*`)**:
|
||||
- **For the LoRA Loader**: Ending a path with `*` acts as a wildcard directory filter. For example, if you input `style/*`, the loader will recursively display all LoRAs located inside any subfolder that begins with `style/`.
|
||||
- **For File Loaders**: If you accidentally copy a path from Windows Explorer that ends in `\*` or `*`, these nodes are smart enough to automatically strip the trailing asterisk so the folder path still resolves cleanly without throwing an OS Error.
|
||||
|
||||
---
|
||||
|
||||
## 4. Settings & Configuration
|
||||
|
||||
Click the standard ComfyUI Settings gear ⚙️ and search for `SaturnNodes` to access the suite's configurations. Here you can:
|
||||
- Safely enter and mask API Keys for Civitai and TMDB scraping.
|
||||
- Change the colors of the Pause Queue buttons via built-in color pickers.
|
||||
- Adjust the default boot behavior for the Persistent Queue (e.g., forcing ComfyUI to always boot in a Paused state), retry limits for failed queue restores, and whether the last manual pause state survives a restart.
|
||||
- Tune the Visual Image Loader metadata scan limit, restore nested output folders, show recorded generation times on restored assets, and trigger **Restore Now** without restarting.
|
||||
- Enable the floating Live Latent dock, drag-resize the queue overlay, and hide entire nodes from the Add Node menu under **Node Management**.
|
||||
|
||||
> **Cross-Platform Note**: All of the visual nodes, pause features, and OS notifications are fully compatible with Windows, macOS, and Linux, and seamlessly adapt to both the modern ComfyUI V2 interface and the classic LiteGraph UI!
|
||||
+578
-54
@@ -1,93 +1,365 @@
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
|
||||
from .nodes.queue_control import PauseQueueNode, PersistentQueueNode, setup_queue_control_routes
|
||||
__version__ = "2.6.3"
|
||||
|
||||
from .nodes.queue_control import (
|
||||
setup_queue_control_routes,
|
||||
pause_manager,
|
||||
tray_manager,
|
||||
persistent_queue_manager,
|
||||
is_persistent_queue_enabled
|
||||
)
|
||||
from .nodes.tray_icon import is_tray_icon_enabled
|
||||
from .nodes.assets_restore import (
|
||||
assets_restore_manager,
|
||||
is_assets_restore_enabled
|
||||
)
|
||||
from .nodes.lora_loader import (
|
||||
FolderLoraLoader,
|
||||
FolderLoraLoaderPretty,
|
||||
FolderLoraLoaderVisualPretty,
|
||||
FolderLoraLoaderVisualPrettyV2
|
||||
VisualLoraLoader
|
||||
)
|
||||
from .nodes.image_loader import ImageLoaderVisualPrettyV2
|
||||
from .nodes.undo_placeholder import UndoPlaceholder
|
||||
from .nodes.image_loader import VisualImageLoader
|
||||
from .nodes.auto_watcher import LoadImageFromFolder
|
||||
from .nodes.load_recent import LoadRecentOutputs
|
||||
from .nodes.preview_latent import PreviewLatentLiveNode
|
||||
from .nodes.decision_node import FlowControlDecision
|
||||
|
||||
from .nodes.decision_node import SaturnDecision
|
||||
from .nodes.aspect_ratio import TextAspectRatioFinder, PreviewImageSizeAspectRatio
|
||||
from .nodes.lora_finder import TextLoraFinder
|
||||
from .nodes.prompt_iterator import PromptQueueIterator
|
||||
from .nodes.prompt_counter import PromptCounter
|
||||
from .nodes.text_replacer import MultiTextReplacer
|
||||
from .nodes.text_split import SaturnTextSplit
|
||||
from .nodes.local_runner import RunLocalFileNode, setup_local_runner_routes
|
||||
from .nodes.stereogram import SaturnStereo3D
|
||||
from .nodes.utils import (
|
||||
get_saturnnodes_user_dir,
|
||||
get_env_setting,
|
||||
is_local_request,
|
||||
is_safe_path,
|
||||
get_csrf_token,
|
||||
is_authenticated_local_request
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FolderLoraLoader": FolderLoraLoader,
|
||||
"FolderLoraLoaderPretty": FolderLoraLoaderPretty,
|
||||
"FolderLoraLoaderVisualPretty": FolderLoraLoaderVisualPretty,
|
||||
"FolderLoraLoaderVisualPrettyV2": FolderLoraLoaderVisualPrettyV2,
|
||||
"ImageLoaderVisualPrettyV2": ImageLoaderVisualPrettyV2,
|
||||
"UndoPlaceholder": UndoPlaceholder,
|
||||
"VisualLoraLoader": VisualLoraLoader,
|
||||
"VisualImageLoader": VisualImageLoader,
|
||||
"LoadImageFromFolder": LoadImageFromFolder,
|
||||
"LoadRecentOutputs": LoadRecentOutputs,
|
||||
"PreviewLatentLive": PreviewLatentLiveNode,
|
||||
"FlowControlDecision": FlowControlDecision
|
||||
"SaturnDecision": SaturnDecision,
|
||||
"TextAspectRatioFinder": TextAspectRatioFinder,
|
||||
"PreviewImageSizeAspectRatio": PreviewImageSizeAspectRatio,
|
||||
"TextLoraFinder": TextLoraFinder,
|
||||
"PromptQueueIterator": PromptQueueIterator,
|
||||
"PromptCounter": PromptCounter,
|
||||
"MultiTextReplacer": MultiTextReplacer,
|
||||
"SaturnTextSplit": SaturnTextSplit,
|
||||
"RunLocalFileNode": RunLocalFileNode,
|
||||
"SaturnStereo3D": SaturnStereo3D,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FolderLoraLoader": "🍃 LoRA Loader (by Folder)",
|
||||
"FolderLoraLoaderPretty": "🍃 LoRA Loader (Pretty)",
|
||||
"FolderLoraLoaderVisualPretty": "🍃 Visual LoRA Picker",
|
||||
"FolderLoraLoaderVisualPrettyV2": "🍃 Visual LoRA Loader",
|
||||
"ImageLoaderVisualPrettyV2": "🍃 Visual Image Loader",
|
||||
"UndoPlaceholder": "🍃 Undo Placeholder",
|
||||
"LoadImageFromFolder": "🍃 Load image from folder",
|
||||
"LoadRecentOutputs": "🍃 Recent Outputs",
|
||||
"PreviewLatentLive": "🍃 Live Latent Preview",
|
||||
"FlowControlDecision": "🍃 FlowControl Decision"
|
||||
"FolderLoraLoader": "🪐 📁 LoRA Loader (Folder)",
|
||||
"FolderLoraLoaderPretty": "🪐 ✨ LoRA Loader (Pretty)",
|
||||
"VisualLoraLoader": "🪐 🖼️ Visual LoRA Loader",
|
||||
"VisualImageLoader": "🪐 📷 Visual Image Loader",
|
||||
"LoadImageFromFolder": "🪐 📂 Load Image From Folder",
|
||||
"LoadRecentOutputs": "🪐 ⏱️ Recent Outputs",
|
||||
"PreviewLatentLive": "🪐 👁️ Live Latent Preview",
|
||||
"SaturnDecision": "🪐 ⏸️ Saturn Decision",
|
||||
"TextAspectRatioFinder": "🪐 📐 Text Aspect Ratio Finder",
|
||||
"PreviewImageSizeAspectRatio": "🪐 📐 Preview Image Size & Aspect Ratio",
|
||||
"TextLoraFinder": "🪐 🔎 Text LoRA Finder & Loader",
|
||||
"PromptQueueIterator": "🪐 🔄 Prompt Queue Iterator",
|
||||
"PromptCounter": "🪐 📝 Prompt Counter",
|
||||
"MultiTextReplacer": "🪐 🔤 Multi Text Replacer",
|
||||
"SaturnTextSplit": "🪐 ✂️ Text Split",
|
||||
"RunLocalFileNode": "🪐 ⚡ Run Local File",
|
||||
"SaturnStereo3D": "🪐 👓 Stereoscopic 3D Generator",
|
||||
}
|
||||
|
||||
ALL_SATURN_NODE_CLASSES = dict(NODE_CLASS_MAPPINGS)
|
||||
ALL_SATURN_NODE_DISPLAY_NAMES = dict(NODE_DISPLAY_NAME_MAPPINGS)
|
||||
ALL_SATURN_NODE_KEYS = list(NODE_CLASS_MAPPINGS)
|
||||
|
||||
def get_saturn_nodes_info():
|
||||
"""Returns sorted node metadata auto-generated from NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS."""
|
||||
import re
|
||||
info_list = []
|
||||
category_order = ["Loaders", "Automation", "Previews", "3D & Previews", "Utils"]
|
||||
|
||||
for key in ALL_SATURN_NODE_KEYS:
|
||||
cls = ALL_SATURN_NODE_CLASSES.get(key)
|
||||
cat = getattr(cls, "CATEGORY", "🪐 SaturnNodes/Utils") if cls else "🪐 SaturnNodes/Utils"
|
||||
subcat = cat.split("/")[-1].strip() if "/" in cat else "Utils"
|
||||
disp = ALL_SATURN_NODE_DISPLAY_NAMES.get(key, key)
|
||||
clean_name = re.sub(r'^[^a-zA-Z0-9\(]+', '', disp).strip()
|
||||
if not clean_name:
|
||||
clean_name = key
|
||||
|
||||
info_list.append({
|
||||
"key": key,
|
||||
"category": subcat,
|
||||
"clean_name": clean_name,
|
||||
"name": f"{subcat} / {clean_name}",
|
||||
"desc": disp
|
||||
})
|
||||
|
||||
def sort_key(item):
|
||||
cat = item["category"]
|
||||
cat_idx = category_order.index(cat) if cat in category_order else 999
|
||||
return (cat_idx, item["name"])
|
||||
|
||||
info_list.sort(key=sort_key)
|
||||
return info_list
|
||||
|
||||
def load_disabled_nodes():
|
||||
"""Loads set of user-disabled SaturnNodes from user settings directory."""
|
||||
try:
|
||||
import json
|
||||
user_dir = get_saturnnodes_user_dir()
|
||||
path = os.path.join(user_dir, "disabled_nodes.json")
|
||||
if os.path.isfile(path):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, list):
|
||||
return set(data)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Notice: Unable to load disabled_nodes.json: {e}")
|
||||
return set()
|
||||
|
||||
def save_disabled_nodes(disabled_set):
|
||||
"""Saves set of user-disabled SaturnNodes atomically."""
|
||||
import json
|
||||
user_dir = get_saturnnodes_user_dir()
|
||||
os.makedirs(user_dir, exist_ok=True)
|
||||
path = os.path.join(user_dir, "disabled_nodes.json")
|
||||
temp_path = path + ".tmp"
|
||||
with open(temp_path, "w", encoding="utf-8") as f:
|
||||
json.dump(sorted(list(disabled_set)), f, indent=2)
|
||||
os.replace(temp_path, path)
|
||||
|
||||
# Filter disabled nodes (all enabled by default when disabled_nodes.json is empty)
|
||||
_active_disabled_nodes = load_disabled_nodes()
|
||||
if _active_disabled_nodes:
|
||||
for _node_name in _active_disabled_nodes:
|
||||
NODE_CLASS_MAPPINGS.pop(_node_name, None)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.pop(_node_name, None)
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
server = PromptServer.instance
|
||||
setup_queue_control_routes(server)
|
||||
server = getattr(PromptServer, "instance", None)
|
||||
if server is not None and hasattr(server, "routes"):
|
||||
setup_queue_control_routes(server)
|
||||
setup_local_runner_routes(server)
|
||||
routes = server.routes
|
||||
else:
|
||||
routes = web.RouteTableDef()
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
ENV_FILE = os.path.join(CURRENT_DIR, ".env")
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
ENV_FILE = os.path.join(USER_DIR, ".env")
|
||||
|
||||
routes = server.routes
|
||||
|
||||
print("[ComfyUI-FlowControl] 🍃 Loaded 11 nodes & visual endpoints successfully.")
|
||||
|
||||
@routes.get("/flow_control/settings")
|
||||
async def get_settings(request):
|
||||
civitai_key = os.getenv("CIVITAI_API_KEY", "")
|
||||
tmdb_key = os.getenv("TMDB_API_KEY", "")
|
||||
if os.path.exists(ENV_FILE):
|
||||
def print_startup_banner():
|
||||
# Enable ANSI escape sequences on Windows console if available
|
||||
if os.name == "nt":
|
||||
try:
|
||||
with open(ENV_FILE, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line.startswith("CIVITAI_API_KEY="):
|
||||
civitai_key = line.split("=", 1)[1].strip()
|
||||
elif line.startswith("TMDB_API_KEY="):
|
||||
tmdb_key = line.split("=", 1)[1].strip()
|
||||
import ctypes
|
||||
kernel32 = ctypes.windll.kernel32
|
||||
kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Saturn Brand Colors (Amber, Gold, Warm Sand, White)
|
||||
AMBER = "\033[1;38;2;245;158;11m"
|
||||
GOLD = "\033[38;2;251;191;36m"
|
||||
SAND = "\033[38;2;222;203;178m"
|
||||
WHITE = "\033[97m"
|
||||
RESET = "\033[0m"
|
||||
|
||||
active_count = len(NODE_DISPLAY_NAME_MAPPINGS)
|
||||
total_count = len(ALL_SATURN_NODE_KEYS)
|
||||
disabled_count = len(_active_disabled_nodes) if _active_disabled_nodes else 0
|
||||
total_str = f"{active_count}/{total_count} nodes loaded"
|
||||
if disabled_count > 0:
|
||||
total_str += f" ({disabled_count} disabled in settings)"
|
||||
|
||||
print()
|
||||
print(f"{AMBER}[ComfyUI-SaturnNodes]{RESET} {SAND}Version:{RESET} {WHITE}{__version__}{RESET}")
|
||||
print(f" {GOLD}*{RESET} {SAND}Loaded:{RESET} {WHITE}Visual LoRA & Image Loaders, LoRA Folder Loader{RESET}")
|
||||
print(f" {GOLD}*{RESET} {SAND}Loaded:{RESET} {WHITE}Live Latent Preview & Floating Dock{RESET}")
|
||||
print(f" {GOLD}*{RESET} {SAND}Loaded:{RESET} {WHITE}Stereoscopic 3D Generator{RESET}")
|
||||
print(f" {GOLD}*{RESET} {SAND}Loaded:{RESET} {WHITE}Prompt Queue Iterator, Prompt Counter, Text Tools{RESET}")
|
||||
print(f" {GOLD}*{RESET} {SAND}Loaded:{RESET} {WHITE}Queue Pause, Persistent Queue, Local Runner{RESET}")
|
||||
print()
|
||||
|
||||
# System status lines
|
||||
try:
|
||||
if is_tray_icon_enabled() and getattr(tray_manager, "_is_running", False):
|
||||
tray_str = "System tray icon started."
|
||||
elif not is_tray_icon_enabled():
|
||||
tray_str = "System tray icon disabled in settings."
|
||||
else:
|
||||
tray_str = "System tray icon not started."
|
||||
except Exception:
|
||||
tray_str = "System tray icon disabled in settings."
|
||||
|
||||
try:
|
||||
if is_assets_restore_enabled():
|
||||
r_count = getattr(assets_restore_manager, "restored_count", 0)
|
||||
assets_str = f"Restored {r_count} recent image(s) into Assets pane."
|
||||
else:
|
||||
assets_str = "No images restored into Assets pane (disabled in settings)."
|
||||
except Exception:
|
||||
assets_str = "Assets pane restore ready."
|
||||
|
||||
try:
|
||||
if is_persistent_queue_enabled():
|
||||
p_items = getattr(persistent_queue_manager, "persistent_items", [])
|
||||
q_count = len(p_items) if p_items else 0
|
||||
queue_str = f"Loaded {q_count} saved queue item(s)."
|
||||
else:
|
||||
queue_str = "No saved queue items loaded (disabled in settings)."
|
||||
except Exception:
|
||||
queue_str = "Persistent queue manager ready."
|
||||
|
||||
print(f" {GOLD}*{RESET} {WHITE}{tray_str}{RESET}")
|
||||
print(f" {GOLD}*{RESET} {WHITE}{assets_str}{RESET}")
|
||||
print(f" {GOLD}*{RESET} {WHITE}{queue_str}{RESET}")
|
||||
print(f"{AMBER}Total: {total_str}{RESET}")
|
||||
print()
|
||||
|
||||
try:
|
||||
print_startup_banner()
|
||||
except Exception:
|
||||
print(f"\n[ComfyUI-SaturnNodes] Version: {__version__} - Total: {len(NODE_DISPLAY_NAME_MAPPINGS)}/{len(ALL_SATURN_NODE_KEYS)} nodes loaded\n")
|
||||
|
||||
@routes.get("/saturnnodes/auth/token")
|
||||
async def get_csrf_token_endpoint(request):
|
||||
if not is_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local access only"}, status=403)
|
||||
return web.json_response({"csrf_token": get_csrf_token()})
|
||||
|
||||
@routes.get("/saturnnodes/get_image_prompt")
|
||||
@routes.get("/saturnnodes/view_image_prompt")
|
||||
async def get_image_prompt_endpoint(request):
|
||||
if not is_local_request(request):
|
||||
return web.json_response({"success": False, "error": "Forbidden: Local access only"}, status=403)
|
||||
try:
|
||||
import folder_paths
|
||||
filename = request.query.get("filename")
|
||||
subfolder = request.query.get("subfolder", "")
|
||||
folder_type = request.query.get("type", "output")
|
||||
|
||||
if not filename:
|
||||
return web.json_response({"success": False, "error": "Missing filename"}, status=400)
|
||||
|
||||
if folder_type == "input":
|
||||
base_dir = folder_paths.get_input_directory()
|
||||
elif folder_type == "temp":
|
||||
base_dir = folder_paths.get_temp_directory()
|
||||
else:
|
||||
base_dir = folder_paths.get_output_directory()
|
||||
|
||||
if subfolder:
|
||||
subfolder_clean = str(subfolder).strip().lstrip("/\\")
|
||||
filepath = os.path.join(base_dir, subfolder_clean, filename)
|
||||
else:
|
||||
filepath = os.path.join(base_dir, filename)
|
||||
|
||||
if not is_safe_path(filepath) or not os.path.exists(filepath) or not os.path.isfile(filepath):
|
||||
return web.json_response({"success": False, "error": "File not found or forbidden"}, status=404)
|
||||
|
||||
from .nodes.image_loader import extract_metadata_from_image
|
||||
prompt, width, height = extract_metadata_from_image(filepath)
|
||||
|
||||
return web.json_response({
|
||||
"success": True,
|
||||
"prompt": prompt,
|
||||
"filename": filename,
|
||||
"width": width,
|
||||
"height": height
|
||||
})
|
||||
except Exception as e:
|
||||
return web.json_response({"success": False, "error": str(e)}, status=500)
|
||||
|
||||
def _env_bool(key, default=False):
|
||||
return str(get_env_setting(key, "true" if default else "false")).lower() in ["true", "1", "yes"]
|
||||
|
||||
def _env_int(key, default):
|
||||
try:
|
||||
return int(get_env_setting(key, str(default)))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
@routes.get("/saturnnodes/settings")
|
||||
async def get_settings(request):
|
||||
if not is_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local access only"}, status=403)
|
||||
# Return the full settings surface (not just API keys) so the settings UI,
|
||||
# debug export and .env never disagree about the effective values.
|
||||
return web.json_response({
|
||||
"civitai_api_key": civitai_key,
|
||||
"tmdb_api_key": tmdb_key
|
||||
"civitai_api_key": get_env_setting("CIVITAI_API_KEY", ""),
|
||||
"tmdb_api_key": get_env_setting("TMDB_API_KEY", ""),
|
||||
"enable_civitai_scraping": _env_bool("ENABLE_CIVITAI_SCRAPING", False),
|
||||
"enable_tmdb_scraping": _env_bool("ENABLE_TMDB_SCRAPING", False),
|
||||
"enable_lora_usage": _env_bool("ENABLE_LORA_USAGE", True),
|
||||
"enable_persistent_queue": _env_bool("ENABLE_PERSISTENT_QUEUE", True),
|
||||
"persistent_queue_restored_state": get_env_setting("PERSISTENT_QUEUE_RESTORED_STATE", "Match Default"),
|
||||
"persistent_queue_max_retries": _env_int("PERSISTENT_QUEUE_MAX_RETRIES", 0),
|
||||
"persist_pause_state": _env_bool("PERSIST_PAUSE_STATE", False),
|
||||
"default_pause_state": get_env_setting("DEFAULT_PAUSE_STATE", "Running"),
|
||||
"default_pause_mode": get_env_setting("DEFAULT_PAUSE_MODE", "after_finish"),
|
||||
"enable_tray_icon": _env_bool("ENABLE_TRAY_ICON", False),
|
||||
"enable_assets_restore": _env_bool("ENABLE_ASSETS_RESTORE", True),
|
||||
"restore_assets_count": _env_int("RESTORE_ASSETS_COUNT", 64),
|
||||
"restore_assets_recursive": _env_bool("RESTORE_ASSETS_RECURSIVE", True),
|
||||
"show_restored_generation_time": _env_bool("SHOW_RESTORED_GENERATION_TIME", True),
|
||||
"image_loader_metadata_limit": _env_int("IMAGE_LOADER_METADATA_LIMIT", 128),
|
||||
"enable_process_management": _env_bool("ALLOW_PROCESS_MANAGEMENT", False),
|
||||
"enable_local_file_execution": _env_bool("ENABLE_LOCAL_FILE_EXECUTION", False),
|
||||
"clear_prompt_iterator_on_launch": _env_bool("CLEAR_PROMPT_ITERATOR_ON_LAUNCH", False)
|
||||
})
|
||||
|
||||
@routes.post("/flow_control/settings")
|
||||
def _clean_env_val(v):
|
||||
if v is None:
|
||||
return None
|
||||
return str(v).replace("\n", "").replace("\r", "").strip()
|
||||
|
||||
@routes.post("/saturnnodes/settings")
|
||||
async def save_settings(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
try:
|
||||
data = await request.json()
|
||||
civitai_key = data.get("civitai_api_key")
|
||||
tmdb_key = data.get("tmdb_api_key")
|
||||
enable_persistent_queue = data.get("enable_persistent_queue")
|
||||
default_pause_state = data.get("default_pause_state")
|
||||
default_pause_mode = data.get("default_pause_mode")
|
||||
enable_civitai = data.get("enable_civitai_scraping")
|
||||
enable_tmdb = data.get("enable_tmdb_scraping")
|
||||
restored_state = data.get("persistent_queue_restored_state")
|
||||
civitai_key = _clean_env_val(data.get("civitai_api_key"))
|
||||
tmdb_key = _clean_env_val(data.get("tmdb_api_key"))
|
||||
enable_persistent_queue = _clean_env_val(data.get("enable_persistent_queue"))
|
||||
default_pause_state = _clean_env_val(data.get("default_pause_state"))
|
||||
default_pause_mode = _clean_env_val(data.get("default_pause_mode"))
|
||||
enable_civitai = _clean_env_val(data.get("enable_civitai_scraping"))
|
||||
enable_tmdb = _clean_env_val(data.get("enable_tmdb_scraping"))
|
||||
enable_lora_usage = _clean_env_val(data.get("enable_lora_usage"))
|
||||
enable_tray_icon = _clean_env_val(data.get("enable_tray_icon"))
|
||||
enable_assets_restore = _clean_env_val(data.get("enable_assets_restore"))
|
||||
restore_assets_count = _clean_env_val(data.get("restore_assets_count"))
|
||||
clear_prompt_iterator_on_launch = _clean_env_val(data.get("clear_prompt_iterator_on_launch"))
|
||||
restored_state = _clean_env_val(data.get("persistent_queue_restored_state"))
|
||||
enable_process_management = _clean_env_val(data.get("enable_process_management"))
|
||||
enable_local_file_execution = _clean_env_val(data.get("enable_local_file_execution"))
|
||||
persistent_queue_max_retries = _clean_env_val(data.get("persistent_queue_max_retries"))
|
||||
persist_pause_state = _clean_env_val(data.get("persist_pause_state"))
|
||||
restore_assets_recursive = _clean_env_val(data.get("restore_assets_recursive"))
|
||||
show_restored_generation_time = _clean_env_val(data.get("show_restored_generation_time"))
|
||||
image_loader_metadata_limit = _clean_env_val(data.get("image_loader_metadata_limit"))
|
||||
|
||||
lines = []
|
||||
if os.path.exists(ENV_FILE):
|
||||
@@ -105,7 +377,19 @@ async def save_settings(request):
|
||||
has_pause_mode = False
|
||||
has_civ_enable = False
|
||||
has_tmdb_enable = False
|
||||
has_lora_usage = False
|
||||
has_tray_enable = False
|
||||
has_assets_restore = False
|
||||
has_restore_count = False
|
||||
has_clear_prompt_iterator = False
|
||||
has_restored_state = False
|
||||
has_process_mgmt = False
|
||||
has_local_file_exec = False
|
||||
has_pq_retries = False
|
||||
has_persist_pause = False
|
||||
has_restore_recursive = False
|
||||
has_show_generation_time = False
|
||||
has_image_metadata_limit = False
|
||||
|
||||
for line in lines:
|
||||
if line.strip().startswith("CIVITAI_API_KEY=") and civitai_key is not None:
|
||||
@@ -129,9 +413,48 @@ async def save_settings(request):
|
||||
elif line.strip().startswith("ENABLE_TMDB_SCRAPING=") and enable_tmdb is not None:
|
||||
new_lines.append(f"ENABLE_TMDB_SCRAPING={enable_tmdb}\n")
|
||||
has_tmdb_enable = True
|
||||
elif line.strip().startswith("ENABLE_LORA_USAGE=") and enable_lora_usage is not None:
|
||||
new_lines.append(f"ENABLE_LORA_USAGE={enable_lora_usage}\n")
|
||||
has_lora_usage = True
|
||||
elif line.strip().startswith("ENABLE_TRAY_ICON=") and enable_tray_icon is not None:
|
||||
new_lines.append(f"ENABLE_TRAY_ICON={enable_tray_icon}\n")
|
||||
has_tray_enable = True
|
||||
elif line.strip().startswith("ENABLE_ASSETS_RESTORE=") and enable_assets_restore is not None:
|
||||
new_lines.append(f"ENABLE_ASSETS_RESTORE={enable_assets_restore}\n")
|
||||
has_assets_restore = True
|
||||
elif line.strip().startswith("RESTORE_ASSETS_COUNT=") and restore_assets_count is not None:
|
||||
new_lines.append(f"RESTORE_ASSETS_COUNT={restore_assets_count}\n")
|
||||
has_restore_count = True
|
||||
elif line.strip().startswith("CLEAR_PROMPT_ITERATOR_ON_LAUNCH=") and clear_prompt_iterator_on_launch is not None:
|
||||
new_lines.append(f"CLEAR_PROMPT_ITERATOR_ON_LAUNCH={clear_prompt_iterator_on_launch}\n")
|
||||
has_clear_prompt_iterator = True
|
||||
elif line.strip().startswith("PERSISTENT_QUEUE_RESTORED_STATE=") and restored_state is not None:
|
||||
new_lines.append(f"PERSISTENT_QUEUE_RESTORED_STATE={restored_state}\n")
|
||||
has_restored_state = True
|
||||
elif line.strip().startswith("ALLOW_PROCESS_MANAGEMENT=") and enable_process_management is not None:
|
||||
new_lines.append(f"ALLOW_PROCESS_MANAGEMENT={enable_process_management}\n")
|
||||
has_process_mgmt = True
|
||||
elif line.strip().startswith("ENABLE_LOCAL_FILE_EXECUTION=") and enable_local_file_execution is not None:
|
||||
new_lines.append(f"ENABLE_LOCAL_FILE_EXECUTION={enable_local_file_execution}\n")
|
||||
has_local_file_exec = True
|
||||
elif line.strip().startswith("PERSISTENT_QUEUE_MAX_RETRIES=") and persistent_queue_max_retries is not None:
|
||||
new_lines.append(f"PERSISTENT_QUEUE_MAX_RETRIES={persistent_queue_max_retries}\n")
|
||||
has_pq_retries = True
|
||||
elif line.strip().startswith("PERSIST_PAUSE_STATE=") and persist_pause_state is not None:
|
||||
new_lines.append(f"PERSIST_PAUSE_STATE={persist_pause_state}\n")
|
||||
has_persist_pause = True
|
||||
elif line.strip().startswith("RESTORE_ASSETS_RECURSIVE=") and restore_assets_recursive is not None:
|
||||
new_lines.append(f"RESTORE_ASSETS_RECURSIVE={restore_assets_recursive}\n")
|
||||
has_restore_recursive = True
|
||||
elif line.strip().startswith("SHOW_RESTORED_GENERATION_TIME=") and show_restored_generation_time is not None:
|
||||
new_lines.append(f"SHOW_RESTORED_GENERATION_TIME={show_restored_generation_time}\n")
|
||||
has_show_generation_time = True
|
||||
elif line.strip().startswith("IMAGE_LOADER_METADATA_LIMIT=") and image_loader_metadata_limit is not None:
|
||||
new_lines.append(f"IMAGE_LOADER_METADATA_LIMIT={image_loader_metadata_limit}\n")
|
||||
has_image_metadata_limit = True
|
||||
elif line.strip().startswith("ALLOW_ANY_SCRIPT_PATH="):
|
||||
# Deprecated and removed for security; drop from .env
|
||||
continue
|
||||
else:
|
||||
new_lines.append(line)
|
||||
|
||||
@@ -149,18 +472,219 @@ async def save_settings(request):
|
||||
new_lines.append(f"ENABLE_CIVITAI_SCRAPING={enable_civitai}\n")
|
||||
if not has_tmdb_enable and enable_tmdb is not None:
|
||||
new_lines.append(f"ENABLE_TMDB_SCRAPING={enable_tmdb}\n")
|
||||
if not has_lora_usage and enable_lora_usage is not None:
|
||||
new_lines.append(f"ENABLE_LORA_USAGE={enable_lora_usage}\n")
|
||||
if not has_tray_enable and enable_tray_icon is not None:
|
||||
new_lines.append(f"ENABLE_TRAY_ICON={enable_tray_icon}\n")
|
||||
if not has_assets_restore and enable_assets_restore is not None:
|
||||
new_lines.append(f"ENABLE_ASSETS_RESTORE={enable_assets_restore}\n")
|
||||
if not has_restore_count and restore_assets_count is not None:
|
||||
new_lines.append(f"RESTORE_ASSETS_COUNT={restore_assets_count}\n")
|
||||
if not has_clear_prompt_iterator and clear_prompt_iterator_on_launch is not None:
|
||||
new_lines.append(f"CLEAR_PROMPT_ITERATOR_ON_LAUNCH={clear_prompt_iterator_on_launch}\n")
|
||||
if not has_restored_state and restored_state is not None:
|
||||
new_lines.append(f"PERSISTENT_QUEUE_RESTORED_STATE={restored_state}\n")
|
||||
if not has_process_mgmt and enable_process_management is not None:
|
||||
new_lines.append(f"ALLOW_PROCESS_MANAGEMENT={enable_process_management}\n")
|
||||
if not has_local_file_exec and enable_local_file_execution is not None:
|
||||
new_lines.append(f"ENABLE_LOCAL_FILE_EXECUTION={enable_local_file_execution}\n")
|
||||
if not has_pq_retries and persistent_queue_max_retries is not None:
|
||||
new_lines.append(f"PERSISTENT_QUEUE_MAX_RETRIES={persistent_queue_max_retries}\n")
|
||||
if not has_persist_pause and persist_pause_state is not None:
|
||||
new_lines.append(f"PERSIST_PAUSE_STATE={persist_pause_state}\n")
|
||||
if not has_restore_recursive and restore_assets_recursive is not None:
|
||||
new_lines.append(f"RESTORE_ASSETS_RECURSIVE={restore_assets_recursive}\n")
|
||||
if not has_show_generation_time and show_restored_generation_time is not None:
|
||||
new_lines.append(f"SHOW_RESTORED_GENERATION_TIME={show_restored_generation_time}\n")
|
||||
if not has_image_metadata_limit and image_loader_metadata_limit is not None:
|
||||
new_lines.append(f"IMAGE_LOADER_METADATA_LIMIT={image_loader_metadata_limit}\n")
|
||||
|
||||
with open(ENV_FILE, "w", encoding="utf-8") as f:
|
||||
temp_env = ENV_FILE + ".tmp"
|
||||
with open(temp_env, "w", encoding="utf-8") as f:
|
||||
f.writelines(new_lines)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(temp_env, ENV_FILE)
|
||||
|
||||
if civitai_key is not None:
|
||||
os.environ["CIVITAI_API_KEY"] = civitai_key
|
||||
if tmdb_key is not None:
|
||||
os.environ["TMDB_API_KEY"] = tmdb_key
|
||||
if enable_lora_usage is not None:
|
||||
os.environ["ENABLE_LORA_USAGE"] = enable_lora_usage
|
||||
if enable_tray_icon is not None:
|
||||
os.environ["ENABLE_TRAY_ICON"] = str(enable_tray_icon)
|
||||
tray_enabled_bool = str(enable_tray_icon).lower() in ["true", "1", "yes"]
|
||||
tray_manager.set_enabled(tray_enabled_bool)
|
||||
if enable_assets_restore is not None:
|
||||
os.environ["ENABLE_ASSETS_RESTORE"] = str(enable_assets_restore)
|
||||
if enable_process_management is not None:
|
||||
os.environ["ALLOW_PROCESS_MANAGEMENT"] = str(enable_process_management)
|
||||
if enable_local_file_execution is not None:
|
||||
os.environ["ENABLE_LOCAL_FILE_EXECUTION"] = str(enable_local_file_execution)
|
||||
if restore_assets_count is not None:
|
||||
os.environ["RESTORE_ASSETS_COUNT"] = str(restore_assets_count)
|
||||
if clear_prompt_iterator_on_launch is not None:
|
||||
os.environ["CLEAR_PROMPT_ITERATOR_ON_LAUNCH"] = str(clear_prompt_iterator_on_launch)
|
||||
if persistent_queue_max_retries is not None:
|
||||
os.environ["PERSISTENT_QUEUE_MAX_RETRIES"] = str(persistent_queue_max_retries)
|
||||
if persist_pause_state is not None:
|
||||
os.environ["PERSIST_PAUSE_STATE"] = str(persist_pause_state)
|
||||
pause_manager.persist_state = str(persist_pause_state).lower() in ["true", "1", "yes"]
|
||||
if pause_manager.persist_state:
|
||||
pause_manager._persist_pause_state()
|
||||
else:
|
||||
try:
|
||||
if os.path.exists(pause_manager.state_file):
|
||||
os.remove(pause_manager.state_file)
|
||||
except OSError:
|
||||
pass
|
||||
if restore_assets_recursive is not None:
|
||||
os.environ["RESTORE_ASSETS_RECURSIVE"] = str(restore_assets_recursive)
|
||||
if show_restored_generation_time is not None:
|
||||
os.environ["SHOW_RESTORED_GENERATION_TIME"] = str(show_restored_generation_time)
|
||||
if image_loader_metadata_limit is not None:
|
||||
os.environ["IMAGE_LOADER_METADATA_LIMIT"] = str(image_loader_metadata_limit)
|
||||
|
||||
# Drop the parsed .env cache immediately so all modules see the new values.
|
||||
from .nodes.utils import invalidate_env_settings_cache
|
||||
invalidate_env_settings_cache()
|
||||
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] 🍃 Error saving settings to .env: {e}")
|
||||
print(f"[SaturnNodes] Error saving settings to .env: {e}")
|
||||
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
||||
|
||||
return web.json_response({"status": "ok"})
|
||||
|
||||
@routes.get("/saturnnodes/disabled_nodes")
|
||||
async def get_disabled_nodes_endpoint(request):
|
||||
if not is_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local access only"}, status=403)
|
||||
return web.json_response({
|
||||
"disabled_nodes": sorted(list(load_disabled_nodes())),
|
||||
"all_nodes": ALL_SATURN_NODE_KEYS,
|
||||
"nodes_info": get_saturn_nodes_info()
|
||||
})
|
||||
|
||||
@routes.post("/saturnnodes/disabled_nodes")
|
||||
async def save_disabled_nodes_endpoint(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
try:
|
||||
data = await request.json()
|
||||
disabled_list = data.get("disabled_nodes", [])
|
||||
if isinstance(disabled_list, list):
|
||||
valid_keys = set(ALL_SATURN_NODE_KEYS)
|
||||
cleaned = {str(k) for k in disabled_list if str(k) in valid_keys}
|
||||
save_disabled_nodes(cleaned)
|
||||
return web.json_response({"status": "ok", "disabled_nodes": sorted(list(cleaned))})
|
||||
return web.json_response({"error": "Invalid format"}, status=400)
|
||||
except Exception as e:
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
@routes.post("/saturnnodes/scrapes/clear")
|
||||
async def clear_scrapes_endpoint(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
try:
|
||||
failed_file = os.path.join(USER_DIR, "failed_scrapes.json")
|
||||
temp_failed = failed_file + ".tmp"
|
||||
with open(temp_failed, "w", encoding="utf-8") as f:
|
||||
f.write("{}")
|
||||
os.replace(temp_failed, failed_file)
|
||||
return web.json_response({"status": "ok"})
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error clearing scrapes cache: {e}")
|
||||
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
||||
|
||||
@routes.get("/saturnnodes/debug/export")
|
||||
async def export_debug_profile(request):
|
||||
if not is_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local access only"}, status=403)
|
||||
try:
|
||||
import platform
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
import torch
|
||||
|
||||
# Safe GPU / hardware information
|
||||
cuda_avail = torch.cuda.is_available() if hasattr(torch, "cuda") else False
|
||||
device_name = torch.cuda.get_device_name(0) if cuda_avail else "CPU"
|
||||
|
||||
# Check API key configuration status without revealing sensitive strings
|
||||
civitai_key = os.getenv("CIVITAI_API_KEY", "")
|
||||
tmdb_key = os.getenv("TMDB_API_KEY", "")
|
||||
if os.path.exists(ENV_FILE):
|
||||
try:
|
||||
with open(ENV_FILE, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if line.strip().startswith("CIVITAI_API_KEY=") and not civitai_key:
|
||||
civitai_key = line.split("=", 1)[1].strip()
|
||||
elif line.strip().startswith("TMDB_API_KEY=") and not tmdb_key:
|
||||
tmdb_key = line.split("=", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Calculate counts of local storage files safely
|
||||
def safe_json_count(filename):
|
||||
p = os.path.join(USER_DIR, filename)
|
||||
if os.path.exists(p):
|
||||
try:
|
||||
with open(p, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return len(data) if isinstance(data, (list, dict)) else 0
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
|
||||
debug_profile = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"package": "ComfyUI-SaturnNodes",
|
||||
"version": __version__,
|
||||
"system": {
|
||||
"os": platform.system(),
|
||||
"os_release": platform.release(),
|
||||
"os_version": platform.version(),
|
||||
"architecture": platform.machine(),
|
||||
"python_version": sys.version.split()[0],
|
||||
"torch_version": getattr(torch, "__version__", "unknown"),
|
||||
"cuda_available": cuda_avail,
|
||||
"device_name": device_name,
|
||||
},
|
||||
"registered_nodes_count": len(NODE_CLASS_MAPPINGS),
|
||||
"settings": {
|
||||
"enable_persistent_queue": get_env_setting("ENABLE_PERSISTENT_QUEUE", "true"),
|
||||
"default_pause_state": get_env_setting("DEFAULT_PAUSE_STATE", "Running"),
|
||||
"default_pause_mode": get_env_setting("DEFAULT_PAUSE_MODE", "after_finish"),
|
||||
"enable_civitai_scraping": get_env_setting("ENABLE_CIVITAI_SCRAPING", "false"),
|
||||
"enable_tmdb_scraping": get_env_setting("ENABLE_TMDB_SCRAPING", "false"),
|
||||
"enable_lora_usage": get_env_setting("ENABLE_LORA_USAGE", "true"),
|
||||
"enable_tray_icon": get_env_setting("ENABLE_TRAY_ICON", "false"),
|
||||
"enable_assets_restore": get_env_setting("ENABLE_ASSETS_RESTORE", "true"),
|
||||
"restore_assets_count": get_env_setting("RESTORE_ASSETS_COUNT", "64"),
|
||||
"restore_assets_recursive": get_env_setting("RESTORE_ASSETS_RECURSIVE", "true"),
|
||||
"show_restored_generation_time": get_env_setting("SHOW_RESTORED_GENERATION_TIME", "true"),
|
||||
"image_loader_metadata_limit": get_env_setting("IMAGE_LOADER_METADATA_LIMIT", "128"),
|
||||
"clear_prompt_iterator_on_launch": get_env_setting("CLEAR_PROMPT_ITERATOR_ON_LAUNCH", "false"),
|
||||
"persistent_queue_restored_state": get_env_setting("PERSISTENT_QUEUE_RESTORED_STATE", "Match Default"),
|
||||
"persistent_queue_max_retries": get_env_setting("PERSISTENT_QUEUE_MAX_RETRIES", "0"),
|
||||
"persist_pause_state": get_env_setting("PERSIST_PAUSE_STATE", "false"),
|
||||
"allow_process_management": get_env_setting("ALLOW_PROCESS_MANAGEMENT", "false"),
|
||||
"enable_local_file_execution": get_env_setting("ENABLE_LOCAL_FILE_EXECUTION", "false"),
|
||||
"civitai_api_key_configured": bool(civitai_key and civitai_key.strip()),
|
||||
"tmdb_api_key_configured": bool(tmdb_key and tmdb_key.strip()),
|
||||
},
|
||||
"local_storage_stats": {
|
||||
"tracked_loras": safe_json_count("lora_usage.json"),
|
||||
"lora_cycle_states": safe_json_count("lora_loader_state.json"),
|
||||
"image_prompts_cached": safe_json_count("image_prompts_cache.json"),
|
||||
"failed_scrapes_cached": safe_json_count("failed_scrapes.json"),
|
||||
"prompt_iterator_queues": safe_json_count("prompt_iterator_state.json"),
|
||||
"persisted_queue_items": safe_json_count("persistent_queue.json"),
|
||||
}
|
||||
}
|
||||
return web.json_response(debug_profile)
|
||||
except Exception as e:
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
|
||||
|
||||
-18
@@ -1,18 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import subprocess
|
||||
|
||||
def main():
|
||||
print("[ComfyUI-FlowControl] Checking dependencies...")
|
||||
dependencies = ["pillow", "piexif", "numpy"]
|
||||
for dep in dependencies:
|
||||
try:
|
||||
__import__(dep)
|
||||
except ImportError:
|
||||
print(f"[ComfyUI-FlowControl] Installing missing dependency: {dep}")
|
||||
subprocess.check_call([sys.executable, "-m", "pip", "install", dep])
|
||||
|
||||
print("[ComfyUI-FlowControl] Installation check complete! Restart ComfyUI to activate.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1 @@
|
||||
# Package marker for SaturnNodes nodes
|
||||
@@ -0,0 +1,253 @@
|
||||
import re
|
||||
import math
|
||||
|
||||
class TextAspectRatioFinder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"aspect_ratios": ("STRING", {
|
||||
"default": "1:1, 2:3, 3:2, 3:4, 4:3, 9:16, 16:9, 21:9",
|
||||
"multiline": False
|
||||
}),
|
||||
"search_mode": (["First match (Front)", "Last match (Back)"], {
|
||||
"default": "First match (Front)"
|
||||
}),
|
||||
"target_mp": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 16.0,
|
||||
"step": 0.1,
|
||||
"display": "number"
|
||||
}),
|
||||
"default_aspect_ratio": ("STRING", {
|
||||
"default": "1:1",
|
||||
"multiline": False
|
||||
}),
|
||||
"multiple_of": ("INT", {
|
||||
"default": 8,
|
||||
"min": 8,
|
||||
"max": 128,
|
||||
"step": 4,
|
||||
"advanced": True
|
||||
}),
|
||||
"min_mp": ("FLOAT", {
|
||||
"default": 0.1,
|
||||
"min": 0.01,
|
||||
"max": 64.0,
|
||||
"step": 0.1,
|
||||
"display": "number",
|
||||
"advanced": True
|
||||
}),
|
||||
"max_mp": ("FLOAT", {
|
||||
"default": 16.0,
|
||||
"min": 0.1,
|
||||
"max": 64.0,
|
||||
"step": 0.1,
|
||||
"display": "number",
|
||||
"advanced": True
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"text": ("STRING", {"forceInput": True}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("width", "height", "aspect_ratio")
|
||||
FUNCTION = "find_aspect_ratio"
|
||||
CATEGORY = "🪐 SaturnNodes/Utils"
|
||||
DESCRIPTION = "Searches text for aspect ratios (e.g. 16:9), syntax checks them, and calculates width & height for target megapixels."
|
||||
|
||||
def find_aspect_ratio(
|
||||
self,
|
||||
aspect_ratios="1:1, 2:3, 3:2, 3:4, 4:3, 9:16, 16:9, 21:9",
|
||||
search_mode="First match (Front)",
|
||||
target_mp=1.0,
|
||||
default_aspect_ratio="1:1",
|
||||
multiple_of=8,
|
||||
min_mp=0.1,
|
||||
max_mp=16.0,
|
||||
text=None,
|
||||
unique_id=None,
|
||||
**kwargs
|
||||
):
|
||||
text_str = str(text) if text is not None else ""
|
||||
|
||||
min_v = float(min_mp)
|
||||
max_v = max(min_v, float(max_mp))
|
||||
effective_mp = max(min_v, min(max_v, float(target_mp)))
|
||||
|
||||
valid_ratios = []
|
||||
raw_items = [item.strip() for item in str(aspect_ratios or "").split(",") if item.strip()]
|
||||
|
||||
for item in raw_items:
|
||||
match = re.match(r'^(\d+(?:\.\d+)?)\s*[:xX/]\s*(\d+(?:\.\d+)?)$', item)
|
||||
if match:
|
||||
w_val = float(match.group(1))
|
||||
h_val = float(match.group(2))
|
||||
if w_val > 0 and h_val > 0:
|
||||
w_str = str(int(w_val)) if w_val.is_integer() else str(w_val)
|
||||
h_str = str(int(h_val)) if h_val.is_integer() else str(h_val)
|
||||
ratio_fmt = f"{w_str}:{h_str}"
|
||||
if ratio_fmt not in valid_ratios:
|
||||
valid_ratios.append(ratio_fmt)
|
||||
|
||||
found_ratio = None
|
||||
matches = []
|
||||
|
||||
if text_str:
|
||||
if valid_ratios:
|
||||
# 1. When allowed aspect_ratios are configured: ONLY match these exact ratios!
|
||||
for ratio in valid_ratios:
|
||||
w_p, h_p = ratio.split(":")
|
||||
pattern = r'(?<![\d.])' + re.escape(w_p) + r'\s*[:xX/]\s*' + re.escape(h_p) + r'(?!\d|\.\d)'
|
||||
for m in re.finditer(pattern, text_str):
|
||||
matches.append((m.start(), ratio))
|
||||
else:
|
||||
# 2. When aspect_ratios input is EMPTY: accept ANY valid aspect ratio found in the text
|
||||
gen_pattern = r'(?<![\d.])(\d+(?:\.\d+)?)\s*[:xX/]\s*(\d+(?:\.\d+)?)(?!\d|\.\d)'
|
||||
for m in re.finditer(gen_pattern, text_str):
|
||||
try:
|
||||
w_val = float(m.group(1))
|
||||
h_val = float(m.group(2))
|
||||
if w_val > 0 and h_val > 0:
|
||||
w_str = str(int(w_val)) if w_val.is_integer() else str(w_val)
|
||||
h_str = str(int(h_val)) if h_val.is_integer() else str(h_val)
|
||||
ratio_fmt = f"{w_str}:{h_str}"
|
||||
if not any(m_pos == m.start() for m_pos, _ in matches):
|
||||
matches.append((m.start(), ratio_fmt))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if matches:
|
||||
matches.sort(key=lambda x: x[0])
|
||||
if search_mode in ["Back", "Last match (Back)"]:
|
||||
found_ratio = matches[-1][1]
|
||||
else:
|
||||
found_ratio = matches[0][1]
|
||||
else:
|
||||
found_ratio = default_aspect_ratio
|
||||
|
||||
# Parse found_ratio and validate. If corrupted/invalid, fall back to "1:1"
|
||||
try:
|
||||
m_ratio = re.match(r'^(\d+(?:\.\d+)?)\s*[:xX/]\s*(\d+(?:\.\d+)?)$', str(found_ratio or "").strip())
|
||||
if m_ratio:
|
||||
w_part = float(m_ratio.group(1))
|
||||
h_part = float(m_ratio.group(2))
|
||||
if w_part <= 0 or h_part <= 0:
|
||||
raise ValueError("Dimensions must be positive")
|
||||
w_str = str(int(w_part)) if w_part.is_integer() else str(w_part)
|
||||
h_str = str(int(h_part)) if h_part.is_integer() else str(h_part)
|
||||
found_ratio = f"{w_str}:{h_str}"
|
||||
else:
|
||||
raise ValueError(f"Cannot parse aspect ratio: {found_ratio}")
|
||||
except Exception as e:
|
||||
err_msg = f"Aspect ratio '{found_ratio}' is invalid. Falling back to '1:1'."
|
||||
print(f"[SaturnNodes] Warning: {err_msg} ({e})")
|
||||
try:
|
||||
from server import PromptServer
|
||||
if hasattr(PromptServer, "instance") and PromptServer.instance:
|
||||
PromptServer.instance.send_sync("saturnnodes_node_error_state", {
|
||||
"node_id": str(unique_id or ""),
|
||||
"title": "Aspect Ratio Warning",
|
||||
"message": err_msg,
|
||||
"fallback": "1:1"
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
w_part, h_part = 1.0, 1.0
|
||||
found_ratio = "1:1"
|
||||
|
||||
r = w_part / h_part
|
||||
|
||||
total_pixels = effective_mp * 1024.0 * 1024.0
|
||||
|
||||
raw_h = math.sqrt(total_pixels / r)
|
||||
raw_w = raw_h * r
|
||||
|
||||
m = max(4, int(multiple_of))
|
||||
width = max(m, int(round(raw_w / m)) * m)
|
||||
height = max(m, int(round(raw_h / m)) * m)
|
||||
|
||||
return (width, height, found_ratio)
|
||||
|
||||
from server import PromptServer
|
||||
|
||||
class PreviewImageSizeAspectRatio:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {},
|
||||
"optional": {
|
||||
"width": ("INT", {"default": 0, "min": 0, "max": 16384, "forceInput": True}),
|
||||
"height": ("INT", {"default": 0, "min": 0, "max": 16384, "forceInput": True}),
|
||||
"aspect_ratio": ("STRING", {"forceInput": True}),
|
||||
"ratio_float": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.01, "forceInput": True}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_NAMES = ()
|
||||
FUNCTION = "process_preview"
|
||||
CATEGORY = "🪐 SaturnNodes/Utils"
|
||||
OUTPUT_NODE = True
|
||||
DESCRIPTION = "Visual preview node that displays image aspect ratio and dimension summary."
|
||||
|
||||
def process_preview(
|
||||
self,
|
||||
width=None,
|
||||
height=None,
|
||||
aspect_ratio=None,
|
||||
ratio_float=None,
|
||||
unique_id="default"
|
||||
):
|
||||
w_val = int(width) if width is not None and int(width) > 0 else 0
|
||||
h_val = int(height) if height is not None and int(height) > 0 else 0
|
||||
ar_str = str(aspect_ratio).strip() if aspect_ratio is not None else ""
|
||||
r_flt = float(ratio_float) if ratio_float is not None and float(ratio_float) > 0 else 0.0
|
||||
|
||||
calc_ratio = 1.0
|
||||
display_text = "1 x 1"
|
||||
|
||||
if w_val > 0 and h_val > 0:
|
||||
calc_ratio = w_val / h_val
|
||||
display_text = f"{w_val} x {h_val}"
|
||||
elif ar_str:
|
||||
match = re.match(r'^(\d+(?:\.\d+)?)\s*[:xX/]\s*(\d+(?:\.\d+)?)$', ar_str)
|
||||
if match:
|
||||
w_p = float(match.group(1))
|
||||
h_p = float(match.group(2))
|
||||
if w_p > 0 and h_p > 0:
|
||||
calc_ratio = w_p / h_p
|
||||
w_disp = int(w_p) if w_p.is_integer() else w_p
|
||||
h_disp = int(h_p) if h_p.is_integer() else h_p
|
||||
display_text = f"{w_disp} x {h_disp}"
|
||||
else:
|
||||
display_text = ar_str
|
||||
elif r_flt > 0:
|
||||
calc_ratio = r_flt
|
||||
if r_flt >= 1.0:
|
||||
w_disp = round(r_flt, 2)
|
||||
display_text = f"{w_disp} x 1"
|
||||
else:
|
||||
h_disp = round(1.0 / r_flt, 2)
|
||||
display_text = f"1 x {h_disp}"
|
||||
|
||||
try:
|
||||
PromptServer.instance.send_sync("saturnnodes_update_preview_aspect_ratio", {
|
||||
"node_id": str(unique_id),
|
||||
"ratio": calc_ratio,
|
||||
"display_text": display_text
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {}
|
||||
|
||||
@@ -0,0 +1,457 @@
|
||||
import os
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
import uuid
|
||||
import heapq
|
||||
import datetime
|
||||
import threading
|
||||
from PIL import Image
|
||||
|
||||
import folder_paths
|
||||
from .utils import get_env_setting
|
||||
from .generation_times import get_generation_time, get_generation_time_debug
|
||||
|
||||
def is_assets_restore_enabled():
|
||||
val = get_env_setting("ENABLE_ASSETS_RESTORE", "true").lower()
|
||||
return val in ["true", "1", "yes"]
|
||||
|
||||
def get_assets_restore_count():
|
||||
try:
|
||||
val = int(get_env_setting("RESTORE_ASSETS_COUNT", "64"))
|
||||
return max(1, min(1000, val))
|
||||
except Exception:
|
||||
return 64
|
||||
|
||||
def is_recursive_assets_scan_enabled():
|
||||
val = get_env_setting("RESTORE_ASSETS_RECURSIVE", "true").lower()
|
||||
return val in ["true", "1", "yes"]
|
||||
|
||||
def is_restored_generation_time_enabled():
|
||||
val = get_env_setting("SHOW_RESTORED_GENERATION_TIME", "true").lower()
|
||||
return val in ["true", "1", "yes"]
|
||||
|
||||
def _strip_exif_text(text):
|
||||
"""Strip NUL padding, a UTF BOM and surrounding whitespace from EXIF text."""
|
||||
value = str(text or "").replace("\x00", "").strip()
|
||||
if value.startswith("\ufeff"):
|
||||
value = value[1:].strip()
|
||||
return value
|
||||
|
||||
def _decode_exif_text(raw):
|
||||
"""Decode an EXIF UserComment, including the standard 8-byte prefix."""
|
||||
if not raw:
|
||||
return ""
|
||||
if isinstance(raw, bytes):
|
||||
prefix, body = raw[:8], raw[8:]
|
||||
if prefix.startswith(b"UNICODE"):
|
||||
# The EXIF standard / piexif write UTF-16 big-endian here, not a
|
||||
# BOM-prefixed native-endian string.
|
||||
try:
|
||||
return _strip_exif_text(body.decode("utf-16-be", errors="ignore"))
|
||||
except Exception:
|
||||
pass
|
||||
if prefix.startswith(b"ASCII"):
|
||||
try:
|
||||
return _strip_exif_text(body.decode("ascii", errors="ignore"))
|
||||
except Exception:
|
||||
pass
|
||||
if prefix.startswith(b"JIS"):
|
||||
try:
|
||||
return _strip_exif_text(body.decode("shift_jis", errors="ignore"))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
return _strip_exif_text(raw.decode("utf-8", errors="ignore"))
|
||||
except Exception:
|
||||
return ""
|
||||
return _strip_exif_text(raw)
|
||||
|
||||
def parse_a1111_positive_prompt(parameters):
|
||||
"""Extract the positive prompt from an A1111/Forge `parameters` string."""
|
||||
if not parameters:
|
||||
return ""
|
||||
if isinstance(parameters, bytes):
|
||||
parameters = _decode_exif_text(parameters)
|
||||
text = str(parameters).replace("\r\n", "\n").replace("\r", "\n").strip()
|
||||
if not text:
|
||||
return ""
|
||||
negative_index = text.find("Negative prompt:")
|
||||
if negative_index >= 0:
|
||||
positive = text[:negative_index].strip()
|
||||
else:
|
||||
positive = text
|
||||
# Drop the trailing generation-settings line (Steps: ..., Sampler: ...).
|
||||
lines = positive.split("\n")
|
||||
if lines:
|
||||
last = lines[-1].strip()
|
||||
if last.startswith("Steps:") or ("Sampler:" in last and "Steps:" in last):
|
||||
lines = lines[:-1]
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
class AssetsRestoreManager:
|
||||
def __init__(self):
|
||||
self._restored_on_launch = False
|
||||
self._restore_lock = threading.RLock()
|
||||
self.last_debug_report = {}
|
||||
self.restored_count = 0
|
||||
|
||||
def extract_image_metadata(self, filepath):
|
||||
"""Extracts embedded prompt/workflow metadata from PNG, JPEG or WebP."""
|
||||
prompt_dict = {}
|
||||
workflow_dict = {}
|
||||
positive_prompt = ""
|
||||
|
||||
try:
|
||||
with Image.open(filepath) as img:
|
||||
prompt_raw = img.info.get("prompt")
|
||||
workflow_raw = img.info.get("workflow")
|
||||
parameters_raw = img.info.get("parameters")
|
||||
|
||||
# JPEG / WebP files written by ComfyUI or A1111 may carry the
|
||||
# metadata in an EXIF UserComment instead of PNG text chunks.
|
||||
if not parameters_raw:
|
||||
exif = None
|
||||
try:
|
||||
exif = img.getexif()
|
||||
except Exception:
|
||||
exif = None
|
||||
if not exif:
|
||||
raw_exif = img.info.get("exif")
|
||||
if raw_exif:
|
||||
try:
|
||||
exif = Image.Exif()
|
||||
exif.load(raw_exif)
|
||||
except Exception:
|
||||
exif = None
|
||||
if exif:
|
||||
try:
|
||||
parameters_raw = _decode_exif_text(exif.get(0x9286))
|
||||
except Exception:
|
||||
parameters_raw = ""
|
||||
|
||||
if prompt_raw:
|
||||
try:
|
||||
parsed_prompt = json.loads(prompt_raw) if isinstance(prompt_raw, str) else prompt_raw
|
||||
if isinstance(parsed_prompt, dict):
|
||||
prompt_dict = parsed_prompt
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if workflow_raw:
|
||||
try:
|
||||
parsed_workflow = json.loads(workflow_raw) if isinstance(workflow_raw, str) else workflow_raw
|
||||
if isinstance(parsed_workflow, dict):
|
||||
workflow_dict = parsed_workflow
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if parameters_raw:
|
||||
positive_prompt = parse_a1111_positive_prompt(parameters_raw)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return prompt_dict, workflow_dict, positive_prompt
|
||||
|
||||
def scan_output_images(self, base_output_dir, limit=64, recursive=True, stats=None):
|
||||
"""Return the newest `limit` images without keeping the full tree in memory."""
|
||||
if not os.path.exists(base_output_dir):
|
||||
return []
|
||||
|
||||
valid_extensions = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
norm_base = os.path.normpath(base_output_dir)
|
||||
if stats is None:
|
||||
stats = {}
|
||||
|
||||
def iter_images():
|
||||
if recursive:
|
||||
roots = os.walk(norm_base)
|
||||
else:
|
||||
try:
|
||||
roots = [(norm_base, [], os.listdir(norm_base))]
|
||||
except OSError:
|
||||
roots = []
|
||||
for root, _, files in roots:
|
||||
for file in files:
|
||||
_, ext = os.path.splitext(file)
|
||||
if ext.lower() not in valid_extensions:
|
||||
continue
|
||||
full_path = os.path.join(root, file)
|
||||
try:
|
||||
mtime = os.path.getmtime(full_path)
|
||||
except OSError:
|
||||
continue
|
||||
stats["scanned"] = stats.get("scanned", 0) + 1
|
||||
rel_dir = os.path.relpath(root, norm_base)
|
||||
subfolder = "" if rel_dir == "." else rel_dir.replace("\\", "/")
|
||||
yield {
|
||||
"filepath": full_path,
|
||||
"filename": file,
|
||||
"subfolder": subfolder,
|
||||
"mtime": mtime,
|
||||
}
|
||||
|
||||
try:
|
||||
safe_limit = max(1, int(limit))
|
||||
except (TypeError, ValueError):
|
||||
safe_limit = 64
|
||||
|
||||
# Keep only the top-N candidates (bounded memory). Include the
|
||||
# normalized relative path in the key so equal mtimes (bulk copies,
|
||||
# archive extraction, FAT timestamp granularity) always resolve to the
|
||||
# same files and the same order on every scan.
|
||||
def sort_key(item):
|
||||
return (
|
||||
item["mtime"],
|
||||
item["subfolder"],
|
||||
item["filename"],
|
||||
)
|
||||
|
||||
top_images = heapq.nlargest(safe_limit, iter_images(), key=sort_key)
|
||||
# Restore chronological order (oldest -> newest) for stable history
|
||||
# numbering. Equal timestamps sort by subfolder/filename ascending.
|
||||
top_images.sort(key=sort_key)
|
||||
return top_images
|
||||
|
||||
def _ingest_assets(self, prepared):
|
||||
"""Register restored files in ComfyUI's asset DB using a single session."""
|
||||
if not prepared:
|
||||
return 0
|
||||
try:
|
||||
import importlib
|
||||
ingest_mod = importlib.import_module("app.assets.services.ingest")
|
||||
register_file_in_place = getattr(ingest_mod, "register_file_in_place", None)
|
||||
if not register_file_in_place:
|
||||
return 0
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
registered = []
|
||||
for entry in prepared:
|
||||
try:
|
||||
res = register_file_in_place(abs_path=entry["filepath"], name=entry["filename"], tags=["output"])
|
||||
if res and hasattr(res, "ref") and res.ref and hasattr(res.ref, "id"):
|
||||
registered.append((res.ref.id, entry))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not registered:
|
||||
return 0
|
||||
|
||||
try:
|
||||
db_mod = importlib.import_module("app.database.db")
|
||||
queries_mod = importlib.import_module("app.assets.database.queries")
|
||||
init_db = getattr(db_mod, "init_db", None)
|
||||
create_session = getattr(db_mod, "create_session", None)
|
||||
set_reference_metadata = getattr(queries_mod, "set_reference_metadata", None)
|
||||
if init_db and create_session and set_reference_metadata:
|
||||
init_db()
|
||||
with create_session() as session:
|
||||
for reference_id, entry in registered:
|
||||
set_reference_metadata(session, reference_id=reference_id, metadata={
|
||||
"jobId": entry["prompt_id"],
|
||||
"nodeId": entry["output_node_id"],
|
||||
"filename": entry["filename"],
|
||||
"subfolder": entry["subfolder"],
|
||||
})
|
||||
session.commit()
|
||||
except Exception:
|
||||
pass
|
||||
return len(registered)
|
||||
|
||||
def restore_on_launch(self, server, limit=None, force=False):
|
||||
"""Populate PromptQueue.history with the newest generated images.
|
||||
|
||||
Serialized because the Restore Now endpoint can be triggered from
|
||||
several browser tabs while the startup restore is still running.
|
||||
"""
|
||||
with self._restore_lock:
|
||||
return self._restore_on_launch(server, limit=limit, force=force)
|
||||
|
||||
def _restore_on_launch(self, server, limit=None, force=False):
|
||||
if self._restored_on_launch and not force:
|
||||
return 0
|
||||
|
||||
started_at = time.time()
|
||||
if limit is None:
|
||||
limit = get_assets_restore_count()
|
||||
else:
|
||||
try:
|
||||
limit = max(1, min(1000, int(limit)))
|
||||
except (TypeError, ValueError):
|
||||
limit = get_assets_restore_count()
|
||||
|
||||
enabled = is_assets_restore_enabled()
|
||||
if not enabled:
|
||||
self._restored_on_launch = True
|
||||
return 0
|
||||
|
||||
recursive = is_recursive_assets_scan_enabled()
|
||||
show_generation_time = is_restored_generation_time_enabled()
|
||||
base_output_dir = folder_paths.get_output_directory()
|
||||
output_exists = os.path.exists(base_output_dir)
|
||||
scan_stats = {}
|
||||
images = self.scan_output_images(
|
||||
base_output_dir, limit=limit, recursive=recursive, stats=scan_stats
|
||||
) if output_exists else []
|
||||
|
||||
prompt_queue = getattr(server, "prompt_queue", None)
|
||||
prepared = []
|
||||
metadata_found = 0
|
||||
generation_time_matches = 0
|
||||
|
||||
# Scan, metadata parsing and history-entry construction happen outside
|
||||
# the queue mutex so the executor thread is never blocked on disk I/O.
|
||||
for idx, item in enumerate(images, 1):
|
||||
filepath = item["filepath"]
|
||||
filename = item["filename"]
|
||||
subfolder = item["subfolder"]
|
||||
|
||||
rel_key = f"{subfolder}/{filename}" if subfolder else filename
|
||||
prompt_id = str(uuid.uuid5(uuid.NAMESPACE_URL, f"saturnnodes:{rel_key}"))
|
||||
|
||||
prompt_dict, workflow_dict, positive_prompt = self.extract_image_metadata(filepath)
|
||||
if prompt_dict or workflow_dict or positive_prompt:
|
||||
metadata_found += 1
|
||||
|
||||
output_node_id = None
|
||||
if isinstance(prompt_dict, dict) and prompt_dict:
|
||||
for nid, nval in prompt_dict.items():
|
||||
if isinstance(nval, dict) and "SaveImage" in str(nval.get("class_type", "")):
|
||||
output_node_id = str(nid)
|
||||
break
|
||||
if not output_node_id:
|
||||
output_node_id = str(next(iter(prompt_dict.keys())))
|
||||
if not output_node_id:
|
||||
output_node_id = "0"
|
||||
|
||||
extra_pnginfo = {}
|
||||
if isinstance(workflow_dict, dict) and workflow_dict:
|
||||
extra_pnginfo["workflow"] = workflow_dict
|
||||
if positive_prompt:
|
||||
extra_pnginfo["positive_prompt"] = positive_prompt
|
||||
|
||||
try:
|
||||
mtime_ms = max(0, int(float(item.get("mtime", 0)) * 1000))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
mtime_ms = 0
|
||||
|
||||
duration = get_generation_time(filename, subfolder, item.get("mtime")) if show_generation_time else None
|
||||
if duration is not None:
|
||||
try:
|
||||
duration = float(duration)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
duration = None
|
||||
if duration is not None and (not math.isfinite(duration) or duration < 0):
|
||||
duration = None
|
||||
|
||||
if duration is not None:
|
||||
generation_time_matches += 1
|
||||
execution_start_ms = max(0, mtime_ms - int(duration * 1000))
|
||||
extra_pnginfo["saturnnodes_generation_time"] = round(duration, 3)
|
||||
else:
|
||||
execution_start_ms = mtime_ms
|
||||
|
||||
# Only emit an execution_success timestamp when a real duration is
|
||||
# known. Two timestamps would make the stock UI report "0.00s" for
|
||||
# older images; start-only preserves their mtime for sorting while
|
||||
# leaving the duration blank.
|
||||
status_messages = [["execution_start", {"timestamp": execution_start_ms}]]
|
||||
if duration is not None:
|
||||
status_messages.append(["execution_success", {"timestamp": mtime_ms}])
|
||||
|
||||
history_entry = {
|
||||
"prompt": [
|
||||
idx,
|
||||
prompt_id,
|
||||
prompt_dict,
|
||||
{
|
||||
"create_time": mtime_ms,
|
||||
"extra_pnginfo": extra_pnginfo,
|
||||
"client_id": "saturnnodes_restored"
|
||||
},
|
||||
[output_node_id]
|
||||
],
|
||||
"outputs": {
|
||||
output_node_id: {
|
||||
"images": [
|
||||
{
|
||||
"filename": filename,
|
||||
"subfolder": subfolder,
|
||||
"type": "output"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"status": {
|
||||
"status_str": "success",
|
||||
"completed": True,
|
||||
"messages": status_messages
|
||||
}
|
||||
}
|
||||
if duration is not None:
|
||||
history_entry["status"]["saturnnodes_duration"] = round(duration, 3)
|
||||
|
||||
prepared.append({
|
||||
"prompt_id": prompt_id,
|
||||
"history_entry": history_entry,
|
||||
"filepath": filepath,
|
||||
"filename": filename,
|
||||
"subfolder": subfolder,
|
||||
"output_node_id": output_node_id,
|
||||
})
|
||||
|
||||
restored_count = 0
|
||||
if prompt_queue is not None:
|
||||
mutex = getattr(prompt_queue, "mutex", None)
|
||||
if mutex is None:
|
||||
mutex = threading.RLock()
|
||||
with mutex:
|
||||
history = getattr(prompt_queue, "history", None)
|
||||
if not isinstance(history, dict):
|
||||
history = {}
|
||||
try:
|
||||
prompt_queue.history = history
|
||||
except Exception:
|
||||
pass
|
||||
for entry in prepared:
|
||||
history[entry["prompt_id"]] = entry["history_entry"]
|
||||
restored_count += 1
|
||||
|
||||
# Asset DB ingestion is best-effort: a broken/absent app.assets must
|
||||
# never prevent the restored history from being usable.
|
||||
try:
|
||||
registered_count = self._ingest_assets(prepared)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Asset DB ingestion failed after restore: {e}")
|
||||
registered_count = 0
|
||||
self._restored_on_launch = True
|
||||
self.restored_count = restored_count
|
||||
|
||||
self.last_debug_report = {
|
||||
"timestamp": datetime.datetime.now().isoformat(),
|
||||
"status": "success",
|
||||
"enabled": enabled,
|
||||
"configured_limit": limit,
|
||||
"recursive": recursive,
|
||||
"show_generation_time": show_generation_time,
|
||||
"output_exists": output_exists,
|
||||
"total_scanned_images": scan_stats.get("scanned", 0),
|
||||
"restored_count": restored_count,
|
||||
"metadata_found": metadata_found,
|
||||
"generation_time_matches": generation_time_matches,
|
||||
"generation_time_records": get_generation_time_debug().get("recorded_files", 0),
|
||||
"asset_db_registered": registered_count,
|
||||
"elapsed_ms": int((time.time() - started_at) * 1000),
|
||||
}
|
||||
|
||||
if restored_count > 0:
|
||||
if server:
|
||||
try:
|
||||
server.queue_updated()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return restored_count
|
||||
|
||||
assets_restore_manager = AssetsRestoreManager()
|
||||
+105
-28
@@ -5,6 +5,39 @@ import numpy as np
|
||||
import re
|
||||
from PIL import Image, ImageOps
|
||||
import comfy.model_management
|
||||
from .utils import sanitize_folder_path, is_safe_path, get_allowed_folder_watch_directories
|
||||
|
||||
def get_current_prompt_id():
|
||||
try:
|
||||
from server import PromptServer
|
||||
server = getattr(PromptServer, "instance", None)
|
||||
if server and hasattr(server, "prompt_queue"):
|
||||
queue = server.prompt_queue
|
||||
if hasattr(queue, "currently_running") and queue.currently_running:
|
||||
mutex = getattr(queue, "mutex", None)
|
||||
if mutex:
|
||||
with mutex:
|
||||
items = list(queue.currently_running.values())
|
||||
else:
|
||||
items = list(queue.currently_running.values())
|
||||
for item in items:
|
||||
if isinstance(item, (tuple, list)) and len(item) > 1:
|
||||
return str(item[1])
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def is_file_ready(filepath):
|
||||
try:
|
||||
if not os.path.exists(filepath) or os.path.getsize(filepath) == 0:
|
||||
return False
|
||||
with open(filepath, "rb") as f:
|
||||
f.seek(0, os.SEEK_END)
|
||||
return True
|
||||
except (OSError, IOError, PermissionError):
|
||||
return False
|
||||
|
||||
_WATCHER_STATE = {}
|
||||
|
||||
class LoadImageFromFolder:
|
||||
@classmethod
|
||||
@@ -14,16 +47,20 @@ class LoadImageFromFolder:
|
||||
"folder": ("STRING", {"default": "input/watch"}),
|
||||
"wait_if_folder_is_empty": ("BOOLEAN", {"default": True}),
|
||||
"rescan_interval": ("INT", {"default": 2, "min": 1, "max": 60}),
|
||||
"sort_by": (["date_modified", "date_created", "name"],),
|
||||
"sort_by": (["date_modified", "date_created", "name"], {"default": "date_modified"}),
|
||||
"regex_filter": ("STRING", {"default": ".*"}),
|
||||
"delete_image": ("BOOLEAN", {"default": False, "tooltip": "If enabled, deletes the image file from disk after loading it. If disabled, cycles through all images in the folder sequentially."}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "BOOLEAN")
|
||||
RETURN_NAMES = ("image", "has_image")
|
||||
FUNCTION = "watch"
|
||||
CATEGORY = "🍃 FlowControl/Automation"
|
||||
DESCRIPTION = "Loads an image from a folder, optionally waiting if the folder is empty."
|
||||
CATEGORY = "🪐 SaturnNodes/Automation"
|
||||
DESCRIPTION = "Loads an image from a folder, optionally waiting if the folder is empty, with automatic sequential cycling when delete is False."
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(s, **kwargs):
|
||||
@@ -31,19 +68,23 @@ class LoadImageFromFolder:
|
||||
return float("NaN")
|
||||
|
||||
def create_dummy_image(self):
|
||||
# Returns a 1x1 0-tensor for no-image state
|
||||
return torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
# Returns a 512x512 standard 0-tensor for safe no-image state
|
||||
return torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
|
||||
def load_and_remove_image(self, filepath):
|
||||
def load_and_remove_image(self, filepath, delete_image=False):
|
||||
with Image.open(filepath) as img:
|
||||
i = ImageOps.exif_transpose(img)
|
||||
image_array = np.array(i.convert("RGB")).astype(np.float32) / 255.0
|
||||
img_tensor = torch.from_numpy(image_array)[None,]
|
||||
|
||||
try:
|
||||
os.remove(filepath)
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Warning: Failed to remove processed file '{filepath}': {e}")
|
||||
if delete_image:
|
||||
if not is_safe_path(filepath, allowed_bases=get_allowed_folder_watch_directories()):
|
||||
print(f"[SaturnNodes Security] Refusing to delete '{filepath}': file is outside ComfyUI input/output directories.")
|
||||
return img_tensor
|
||||
try:
|
||||
os.remove(filepath)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Warning: Failed to remove processed file '{filepath}': {e}")
|
||||
|
||||
return img_tensor
|
||||
|
||||
@@ -52,14 +93,14 @@ class LoadImageFromFolder:
|
||||
try:
|
||||
regex = re.compile(regex_filter)
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Invalid regex: {e}. Falling back to .*")
|
||||
print(f"[SaturnNodes] Invalid regex: {e}. Falling back to .*")
|
||||
regex = re.compile(".*")
|
||||
|
||||
files = []
|
||||
for f in os.listdir(folder):
|
||||
filepath = os.path.join(folder, f)
|
||||
if os.path.isfile(filepath) and f.lower().endswith(valid_extensions):
|
||||
if regex.search(f):
|
||||
if regex.search(f) and is_file_ready(filepath):
|
||||
files.append(filepath)
|
||||
|
||||
if not files:
|
||||
@@ -85,7 +126,41 @@ class LoadImageFromFolder:
|
||||
|
||||
return [x[0] for x in valid_files]
|
||||
|
||||
def watch(self, folder, wait_if_folder_is_empty, rescan_interval, sort_by, regex_filter):
|
||||
def _select_file(self, files, delete_image, unique_id, folder):
|
||||
if not files:
|
||||
return None
|
||||
|
||||
if delete_image:
|
||||
# When deleting on each run, always consume the first available file
|
||||
return files[0]
|
||||
|
||||
# When keeping images (delete_image=False), cycle sequentially through all files!
|
||||
state_key = f"{unique_id}_{folder}"
|
||||
current_pid = get_current_prompt_id()
|
||||
|
||||
node_state = _WATCHER_STATE.get(state_key, {"index": 0, "assigned": {}})
|
||||
assigned = node_state.get("assigned", {})
|
||||
current_idx = int(node_state.get("index", 0))
|
||||
|
||||
if current_pid and current_pid in assigned:
|
||||
# Re-use exact index on persistent queue crash recovery
|
||||
idx = assigned[current_pid] % len(files)
|
||||
else:
|
||||
idx = current_idx % len(files)
|
||||
if current_pid:
|
||||
assigned[current_pid] = idx
|
||||
if len(assigned) > 50:
|
||||
for k in list(assigned.keys())[:-50]:
|
||||
del assigned[k]
|
||||
node_state["index"] = (current_idx + 1) % len(files)
|
||||
node_state["assigned"] = assigned
|
||||
_WATCHER_STATE[state_key] = node_state
|
||||
|
||||
return files[idx]
|
||||
|
||||
def watch(self, folder, wait_if_folder_is_empty, rescan_interval, sort_by, regex_filter, delete_image=False, unique_id="default", **kwargs):
|
||||
folder = sanitize_folder_path(folder, default_dir="watch")
|
||||
|
||||
if wait_if_folder_is_empty:
|
||||
while True:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
@@ -94,14 +169,15 @@ class LoadImageFromFolder:
|
||||
try:
|
||||
files = self.get_filtered_files(folder, sort_by, regex_filter)
|
||||
if files:
|
||||
filepath = files[0] # Pick the first one based on sorting
|
||||
try:
|
||||
img_tensor = self.load_and_remove_image(filepath)
|
||||
return (img_tensor, True)
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Error processing {filepath}: {e}")
|
||||
filepath = self._select_file(files, delete_image, unique_id, folder)
|
||||
if filepath:
|
||||
try:
|
||||
img_tensor = self.load_and_remove_image(filepath, delete_image=delete_image)
|
||||
return (img_tensor, True)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error processing {filepath}: {e}")
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Error accessing directory {folder}: {e}")
|
||||
print(f"[SaturnNodes] Error accessing directory {folder}: {e}")
|
||||
|
||||
time.sleep(rescan_interval)
|
||||
else:
|
||||
@@ -109,14 +185,15 @@ class LoadImageFromFolder:
|
||||
try:
|
||||
files = self.get_filtered_files(folder, sort_by, regex_filter)
|
||||
if files:
|
||||
filepath = files[0]
|
||||
try:
|
||||
img_tensor = self.load_and_remove_image(filepath)
|
||||
return (img_tensor, True)
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Error processing {filepath}: {e}")
|
||||
filepath = self._select_file(files, delete_image, unique_id, folder)
|
||||
if filepath:
|
||||
try:
|
||||
img_tensor = self.load_and_remove_image(filepath, delete_image=delete_image)
|
||||
return (img_tensor, True)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error processing {filepath}: {e}")
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Error accessing directory {folder}: {e}")
|
||||
print(f"[SaturnNodes] Error accessing directory {folder}: {e}")
|
||||
|
||||
# No image present or folder missing -> return dummy tensor & False immediately
|
||||
# No image present or folder missing -> return standard dummy tensor & False immediately
|
||||
return (self.create_dummy_image(), False)
|
||||
|
||||
+81
-37
@@ -1,9 +1,11 @@
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import subprocess
|
||||
from server import PromptServer
|
||||
from aiohttp import web
|
||||
import comfy.model_management
|
||||
from .utils import is_authenticated_local_request
|
||||
|
||||
class DecisionManager:
|
||||
_waiting_events = {}
|
||||
@@ -36,13 +38,13 @@ class DecisionManager:
|
||||
return False
|
||||
|
||||
|
||||
class FlowControlDecision:
|
||||
class SaturnDecision:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"disable": ("BOOLEAN", {"default": False}),
|
||||
"send_windows_notification": ("BOOLEAN", {"default": False}),
|
||||
"send_os_notification": ("BOOLEAN", {"default": False}),
|
||||
"timeout": ("INT", {"default": -1, "min": -1, "max": 3600}),
|
||||
},
|
||||
"hidden": {
|
||||
@@ -53,13 +55,16 @@ class FlowControlDecision:
|
||||
RETURN_TYPES = ("BOOLEAN",)
|
||||
RETURN_NAMES = ("cancel",)
|
||||
FUNCTION = "decide"
|
||||
CATEGORY = "🍃 FlowControl"
|
||||
DESCRIPTION = "Pauses execution and waits for your input via the UI buttons.\n\n- 'Continue' outputs False (0).\n- 'Cancel' outputs True (1) so you can route it into a Switch node to bypass later nodes.\n- 'Stop Workflow' instantly aborts the entire ComfyUI generation queue."
|
||||
CATEGORY = "🪐 SaturnNodes/Utils"
|
||||
DESCRIPTION = "Pauses execution and waits for your input via the UI buttons.\n\n- 'Continue' outputs False (0).\n- 'Cancel' outputs True (1) so you can route it into a Switch node to bypass later nodes.\n- 'Stop Workflow' instantly aborts the entire ComfyUI generation queue.\n- 'OS Notification': Sends a native desktop toast (Windows/macOS/Linux) when waiting."
|
||||
|
||||
def send_notification(self, title, message):
|
||||
if os.name == 'nt':
|
||||
try:
|
||||
# Using native PowerShell Toast Notification API (no third party modules required)
|
||||
try:
|
||||
if sys.platform == 'win32' or os.name == 'nt':
|
||||
# Windows native PowerShell Toast Notification with XML escaping
|
||||
import xml.sax.saxutils
|
||||
esc_title = xml.sax.saxutils.escape(str(title))
|
||||
esc_msg = xml.sax.saxutils.escape(str(message))
|
||||
ps_script = f"""
|
||||
$ErrorActionPreference = 'Stop'
|
||||
[Windows.UI.Notifications.ToastNotificationManager, Windows.UI.Notifications, ContentType = WindowsRuntime] | Out-Null
|
||||
@@ -69,8 +74,8 @@ $template = @"
|
||||
<toast>
|
||||
<visual>
|
||||
<binding template="ToastText02">
|
||||
<text id="1">{title}</text>
|
||||
<text id="2">{message}</text>
|
||||
<text id="1">{esc_title}</text>
|
||||
<text id="2">{esc_msg}</text>
|
||||
</binding>
|
||||
</visual>
|
||||
</toast>
|
||||
@@ -82,10 +87,19 @@ $toast = [Windows.UI.Notifications.ToastNotification]::new($xml)
|
||||
[Windows.UI.Notifications.ToastNotificationManager]::CreateToastNotifier("ComfyUI").Show($toast)
|
||||
"""
|
||||
subprocess.Popen(["powershell", "-NoProfile", "-Command", ps_script], creationflags=subprocess.CREATE_NO_WINDOW)
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Failed to send Windows notification: {e}")
|
||||
elif sys.platform == 'darwin':
|
||||
# macOS native AppleScript Notification
|
||||
escape_title = title.replace('"', '\\"')
|
||||
escape_message = message.replace('"', '\\"')
|
||||
apple_script = f'display notification "{escape_message}" with title "{escape_title}"'
|
||||
subprocess.Popen(["osascript", "-e", apple_script])
|
||||
else:
|
||||
# Linux dbus notify-send
|
||||
subprocess.Popen(["notify-send", title, message])
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Failed to send OS notification: {e}")
|
||||
|
||||
def decide(self, disable, send_windows_notification, timeout, unique_id=None):
|
||||
def decide(self, disable, send_os_notification, timeout, unique_id=None):
|
||||
if disable:
|
||||
return (False,)
|
||||
|
||||
@@ -93,57 +107,87 @@ $toast = [Windows.UI.Notifications.ToastNotification]::new($xml)
|
||||
return (False,)
|
||||
|
||||
# Notify frontend
|
||||
PromptServer.instance.send_sync("flowcontrol_decision_waiting", {"node_id": unique_id})
|
||||
PromptServer.instance.send_sync("saturnnodes_decision_waiting", {"node_id": unique_id})
|
||||
|
||||
if send_windows_notification:
|
||||
self.send_notification("ComfyUI FlowControl", "Workflow paused! Waiting for your decision.")
|
||||
if send_os_notification:
|
||||
self.send_notification("ComfyUI SaturnNodes", "Workflow paused! Waiting for your decision.")
|
||||
|
||||
# Wait for user input
|
||||
event = threading.Event()
|
||||
DecisionManager.register_wait(unique_id, event)
|
||||
|
||||
print(f"[FlowControl] Node {unique_id} is waiting for user decision...")
|
||||
|
||||
if timeout > 0:
|
||||
waited = event.wait(timeout)
|
||||
if not waited:
|
||||
print(f"[FlowControl] Node {unique_id} timed out. Auto-continuing...")
|
||||
DecisionManager.unregister_wait(unique_id)
|
||||
# Notify frontend to update UI back to normal
|
||||
PromptServer.instance.send_sync("flowcontrol_decision_resolved", {"node_id": unique_id})
|
||||
return (False,)
|
||||
else:
|
||||
event.wait()
|
||||
import time
|
||||
try:
|
||||
if timeout > 0:
|
||||
end_time = time.time() + timeout
|
||||
while time.time() < end_time and not event.is_set():
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
event.wait(0.5)
|
||||
if not event.is_set():
|
||||
print(f"[SaturnNodes] Node {unique_id} timed out. Auto-continuing...")
|
||||
DecisionManager.unregister_wait(unique_id)
|
||||
# Notify frontend to update UI back to normal
|
||||
PromptServer.instance.send_sync("saturnnodes_decision_resolved", {"node_id": unique_id})
|
||||
return (False,)
|
||||
else:
|
||||
while not event.is_set():
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
event.wait(0.5)
|
||||
except Exception:
|
||||
DecisionManager.unregister_wait(unique_id)
|
||||
PromptServer.instance.send_sync("saturnnodes_decision_resolved", {"node_id": unique_id})
|
||||
raise
|
||||
|
||||
# Retrieve action
|
||||
action = DecisionManager.get_action(unique_id)
|
||||
|
||||
# Notify frontend that the decision has been resolved
|
||||
PromptServer.instance.send_sync("flowcontrol_decision_resolved", {"node_id": unique_id})
|
||||
PromptServer.instance.send_sync("saturnnodes_decision_resolved", {"node_id": unique_id})
|
||||
|
||||
if action == "cancel":
|
||||
print(f"[FlowControl] Node {unique_id} cancelled.")
|
||||
print(f"[SaturnNodes] Node {unique_id} cancelled.")
|
||||
return (True,)
|
||||
elif action == "stop":
|
||||
print(f"[FlowControl] Node {unique_id} stopped workflow.")
|
||||
print(f"[SaturnNodes] Node {unique_id} stopped workflow.")
|
||||
stop_dispatched = False
|
||||
try:
|
||||
# Attempt to cancel current execution queue
|
||||
PromptServer.instance.prompt_queue.cancel_current_execution()
|
||||
except:
|
||||
pass
|
||||
if hasattr(PromptServer, "instance") and hasattr(PromptServer.instance, "prompt_queue"):
|
||||
PromptServer.instance.prompt_queue.cancel_current_execution()
|
||||
stop_dispatched = True
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Notice: cancel_current_execution: {e}")
|
||||
try:
|
||||
import nodes
|
||||
if hasattr(nodes, "interrupt_processing"):
|
||||
nodes.interrupt_processing()
|
||||
stop_dispatched = True
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Notice: nodes.interrupt_processing: {e}")
|
||||
try:
|
||||
comfy.model_management.interrupt_current_processing()
|
||||
stop_dispatched = True
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Notice: interrupt_current_processing: {e}")
|
||||
|
||||
if not stop_dispatched:
|
||||
print(f"[SaturnNodes] Notice: Core cancellation hooks bypassed; halting via InterruptProcessingException.")
|
||||
|
||||
# Raise exception to halt this thread instantly
|
||||
try:
|
||||
raise comfy.model_management.InterruptProcessingException("Workflow stopped by FlowControl Decision Node.")
|
||||
raise comfy.model_management.InterruptProcessingException("Workflow stopped by SaturnNodes Decision Node.")
|
||||
except AttributeError:
|
||||
raise Exception("Workflow stopped by FlowControl Decision Node.")
|
||||
raise RuntimeError("Workflow stopped by SaturnNodes Decision Node.")
|
||||
|
||||
print(f"[FlowControl] Node {unique_id} continuing.")
|
||||
print(f"[SaturnNodes] Node {unique_id} continuing.")
|
||||
return (False,)
|
||||
|
||||
|
||||
# Register API Route
|
||||
@PromptServer.instance.routes.post("/flowcontrol/decision")
|
||||
@PromptServer.instance.routes.post("/saturnnodes/decision")
|
||||
async def handle_decision(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
try:
|
||||
data = await request.json()
|
||||
node_id = data.get("node_id")
|
||||
|
||||
@@ -0,0 +1,329 @@
|
||||
"""
|
||||
Generation-time tracking for SaturnNodes asset restoration.
|
||||
|
||||
ComfyUI measures prompt execution duration in main.py but only logs it; it is
|
||||
not persisted in history or PNG metadata. This module wraps PromptQueue.get /
|
||||
task_done once and records duration + output filenames in a small JSON file so
|
||||
restored history entries can show a real execution time.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
import threading
|
||||
from .utils import get_saturnnodes_user_dir
|
||||
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
GENERATION_TIMES_FILE = os.path.join(USER_DIR, "generation_times.json")
|
||||
MAX_ENTRIES = 1000
|
||||
|
||||
_LOCK = threading.RLock()
|
||||
_DATA = None
|
||||
_DATA_MTIME = None
|
||||
_START_TIMES = {}
|
||||
_PATCHED = False
|
||||
|
||||
|
||||
def _entry_timestamp(entry):
|
||||
if not isinstance(entry, dict):
|
||||
return 0
|
||||
try:
|
||||
value = int(entry.get("completed", 0))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return 0
|
||||
return value if value >= 0 else 0
|
||||
|
||||
|
||||
def _sanitize_data(data):
|
||||
"""Return a JSON-safe copy of the duration store.
|
||||
|
||||
Hand-edited or partially written files can contain non-finite, negative or
|
||||
non-numeric durations. Dropping them here keeps one corrupt entry from
|
||||
making every future save fail under ``allow_nan=False``.
|
||||
"""
|
||||
payload = dict(data) if isinstance(data, dict) else {}
|
||||
payload.setdefault("version", 1)
|
||||
entries = {}
|
||||
raw_entries = payload.get("entries")
|
||||
if isinstance(raw_entries, dict):
|
||||
for key, entry in raw_entries.items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
try:
|
||||
duration = float(entry.get("duration"))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
continue
|
||||
if not math.isfinite(duration) or duration < 0:
|
||||
continue
|
||||
cleaned = dict(entry)
|
||||
cleaned["duration"] = duration
|
||||
try:
|
||||
completed = int(entry.get("completed", 0))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
completed = 0
|
||||
cleaned["completed"] = completed if completed >= 0 else 0
|
||||
entries[key] = cleaned
|
||||
payload["entries"] = entries
|
||||
return payload
|
||||
|
||||
|
||||
def _load_data():
|
||||
global _DATA, _DATA_MTIME
|
||||
with _LOCK:
|
||||
mtime = None
|
||||
try:
|
||||
if os.path.exists(GENERATION_TIMES_FILE):
|
||||
mtime = os.stat(GENERATION_TIMES_FILE).st_mtime_ns
|
||||
except OSError:
|
||||
mtime = None
|
||||
if _DATA is not None and mtime == _DATA_MTIME:
|
||||
return _DATA
|
||||
|
||||
data = {"version": 1, "entries": {}}
|
||||
if mtime is not None:
|
||||
try:
|
||||
with open(GENERATION_TIMES_FILE, "r", encoding="utf-8") as f:
|
||||
loaded = json.load(f)
|
||||
if isinstance(loaded, dict) and isinstance(loaded.get("entries"), dict):
|
||||
data = _sanitize_data(loaded)
|
||||
except Exception:
|
||||
pass
|
||||
_DATA = data
|
||||
_DATA_MTIME = mtime
|
||||
return _DATA
|
||||
|
||||
|
||||
def _save_data(data):
|
||||
global _DATA_MTIME
|
||||
temp_file = None
|
||||
try:
|
||||
os.makedirs(USER_DIR, exist_ok=True)
|
||||
data = _sanitize_data(data)
|
||||
# Unique temp names keep a stale temp file from a crashed process (or a
|
||||
# second ComfyUI instance) from being clobbered mid-write.
|
||||
temp_file = f"{GENERATION_TIMES_FILE}.{os.getpid()}.{threading.get_ident()}.tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, separators=(",", ":"), ensure_ascii=False, allow_nan=False)
|
||||
f.flush()
|
||||
try:
|
||||
os.fsync(f.fileno())
|
||||
except OSError:
|
||||
pass
|
||||
for attempt in range(3):
|
||||
try:
|
||||
os.replace(temp_file, GENERATION_TIMES_FILE)
|
||||
temp_file = None
|
||||
break
|
||||
except PermissionError:
|
||||
if attempt == 2:
|
||||
raise
|
||||
time.sleep(0.05)
|
||||
try:
|
||||
_DATA_MTIME = os.stat(GENERATION_TIMES_FILE).st_mtime_ns
|
||||
except OSError:
|
||||
_DATA_MTIME = None
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error saving generation times: {e}")
|
||||
finally:
|
||||
if temp_file and os.path.exists(temp_file):
|
||||
try:
|
||||
os.remove(temp_file)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _normalize_ref(filename, subfolder=""):
|
||||
sub = str(subfolder or "").replace("\\", "/").strip("/")
|
||||
name = str(filename or "").replace("\\", "/").strip("/")
|
||||
return f"{sub}/{name}" if sub else name
|
||||
|
||||
|
||||
def record_generation_time(prompt_id, duration, files, status="success"):
|
||||
"""Store one prompt's wall-clock duration keyed by every output file."""
|
||||
try:
|
||||
duration = round(float(duration), 3)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return
|
||||
if not math.isfinite(duration) or duration < 0 or not files:
|
||||
return
|
||||
|
||||
completed_ms = int(time.time() * 1000)
|
||||
with _LOCK:
|
||||
data = _load_data()
|
||||
entries = data.setdefault("entries", {})
|
||||
for file_info in files:
|
||||
if not isinstance(file_info, dict):
|
||||
continue
|
||||
filename = file_info.get("filename")
|
||||
if not filename:
|
||||
continue
|
||||
key = _normalize_ref(filename, file_info.get("subfolder", ""))
|
||||
entries[key] = {
|
||||
"duration": duration,
|
||||
"completed": completed_ms,
|
||||
"prompt_id": str(prompt_id) if prompt_id else "",
|
||||
"status": str(status) if status is not None else "success",
|
||||
}
|
||||
|
||||
if len(entries) > MAX_ENTRIES:
|
||||
keep = set(sorted(entries.keys(), key=lambda k: _entry_timestamp(entries[k]), reverse=True)[:MAX_ENTRIES])
|
||||
for key in [k for k in entries.keys() if k not in keep]:
|
||||
del entries[key]
|
||||
|
||||
_save_data(data)
|
||||
|
||||
|
||||
def get_generation_time(filename, subfolder="", file_mtime=None):
|
||||
"""Return the recorded duration in seconds for an output file, or None.
|
||||
|
||||
``file_mtime`` (seconds since epoch) is optional and only used for the
|
||||
moved-file fallback. It prevents two different files that happen to share a
|
||||
bare filename from being assigned each other's duration.
|
||||
"""
|
||||
if not filename:
|
||||
return None
|
||||
key = _normalize_ref(filename, subfolder)
|
||||
with _LOCK:
|
||||
entries = _load_data().get("entries", {})
|
||||
entry = entries.get(key)
|
||||
if entry is None:
|
||||
bare = str(filename).replace("\\", "/").strip("/")
|
||||
matches = [
|
||||
(k, v) for k, v in entries.items()
|
||||
if k == bare or k.endswith("/" + bare)
|
||||
]
|
||||
if file_mtime is not None:
|
||||
try:
|
||||
mtime_ms = float(file_mtime) * 1000.0
|
||||
matches = [
|
||||
(k, v) for k, v in matches
|
||||
if abs(_entry_timestamp(v) - mtime_ms) <= 300000
|
||||
]
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
# Only use the fallback when exactly one record can belong to this
|
||||
# file. Guessing between same-named outputs is worse than showing
|
||||
# no duration at all.
|
||||
if len(matches) == 1:
|
||||
entry = matches[0][1]
|
||||
if isinstance(entry, dict):
|
||||
try:
|
||||
duration = float(entry.get("duration"))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return None
|
||||
if math.isfinite(duration) and duration >= 0:
|
||||
return duration
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def get_generation_time_debug():
|
||||
with _LOCK:
|
||||
entries = _load_data().get("entries", {})
|
||||
return {"recorded_files": len(entries), "max_entries": MAX_ENTRIES}
|
||||
|
||||
|
||||
def _extract_output_files(history_result):
|
||||
files = []
|
||||
if not isinstance(history_result, dict):
|
||||
return files
|
||||
outputs = history_result.get("outputs")
|
||||
if not isinstance(outputs, dict):
|
||||
return files
|
||||
for node_output in outputs.values():
|
||||
if not isinstance(node_output, dict):
|
||||
continue
|
||||
for image in node_output.get("images", []) or []:
|
||||
if not isinstance(image, dict) or not image.get("filename"):
|
||||
continue
|
||||
# Only permanently saved outputs are eligible for restored-asset
|
||||
# history entries. Temporary previews can share filenames.
|
||||
if image.get("type", "output") != "output":
|
||||
continue
|
||||
files.append({
|
||||
"filename": image.get("filename"),
|
||||
"subfolder": image.get("subfolder", ""),
|
||||
"type": "output",
|
||||
})
|
||||
return files
|
||||
|
||||
|
||||
def patch_prompt_queue(server):
|
||||
"""Wrap PromptQueue.get/task_done once to measure real execution durations."""
|
||||
global _PATCHED
|
||||
if _PATCHED:
|
||||
return
|
||||
queue = getattr(server, "prompt_queue", None)
|
||||
if queue is None or not hasattr(queue, "get") or not hasattr(queue, "task_done"):
|
||||
return
|
||||
if getattr(queue, "_saturnnodes_time_patched", False):
|
||||
_PATCHED = True
|
||||
return
|
||||
|
||||
original_get = queue.get
|
||||
|
||||
def timed_get(*args, **kwargs):
|
||||
item = original_get(*args, **kwargs)
|
||||
if item is not None:
|
||||
try:
|
||||
# ComfyUI's PromptQueue.get returns (queue_item, task_id).
|
||||
# Support a bare queue item as well for lightweight/legacy
|
||||
# queue implementations used by tests or custom servers.
|
||||
queue_item = item[0] if (
|
||||
isinstance(item, (tuple, list))
|
||||
and len(item) == 2
|
||||
and isinstance(item[0], (tuple, list))
|
||||
) else item
|
||||
prompt_id = str(queue_item[1]) if (
|
||||
isinstance(queue_item, (tuple, list)) and len(queue_item) > 1
|
||||
) else None
|
||||
if prompt_id:
|
||||
with _LOCK:
|
||||
_START_TIMES[prompt_id] = time.time()
|
||||
while len(_START_TIMES) > 2000:
|
||||
_START_TIMES.pop(next(iter(_START_TIMES)))
|
||||
except Exception:
|
||||
pass
|
||||
return item
|
||||
|
||||
original_task_done = queue.task_done
|
||||
|
||||
def timed_task_done(item_id, history_result, status=None, process_item=None):
|
||||
prompt_id = None
|
||||
try:
|
||||
mutex = getattr(queue, "mutex", None)
|
||||
if mutex:
|
||||
with mutex:
|
||||
running = queue.currently_running.get(item_id)
|
||||
else:
|
||||
running = queue.currently_running.get(item_id)
|
||||
if isinstance(running, (tuple, list)) and len(running) > 1:
|
||||
prompt_id = str(running[1])
|
||||
except Exception:
|
||||
prompt_id = None
|
||||
|
||||
started = None
|
||||
with _LOCK:
|
||||
if prompt_id:
|
||||
started = _START_TIMES.pop(prompt_id, None)
|
||||
|
||||
try:
|
||||
return original_task_done(item_id, history_result, status, process_item)
|
||||
finally:
|
||||
if prompt_id and started is not None:
|
||||
status_str = getattr(status, "status_str", status)
|
||||
record_generation_time(
|
||||
prompt_id,
|
||||
max(0.0, time.time() - started),
|
||||
_extract_output_files(history_result),
|
||||
status=status_str,
|
||||
)
|
||||
|
||||
queue.get = timed_get
|
||||
queue.task_done = timed_task_done
|
||||
try:
|
||||
queue._saturnnodes_time_patched = True
|
||||
except Exception:
|
||||
pass
|
||||
_PATCHED = True
|
||||
+312
-80
@@ -4,19 +4,44 @@ import torch
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps
|
||||
import io
|
||||
import asyncio
|
||||
import hashlib
|
||||
import threading
|
||||
import piexif
|
||||
import piexif.helper
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
import folder_paths
|
||||
from .utils import (
|
||||
sanitize_folder_path,
|
||||
is_local_request,
|
||||
is_safe_path,
|
||||
sanitize_image_loader_folder,
|
||||
get_saturnnodes_user_dir,
|
||||
get_env_setting
|
||||
)
|
||||
|
||||
IMAGE_CATEGORY = "🍃 FlowControl/Loaders"
|
||||
IMAGE_CATEGORY = "🪐 SaturnNodes/Loaders"
|
||||
CURRENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
CACHE_FILE = os.path.join(CURRENT_DIR, "image_prompts_cache.json")
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
CACHE_FILE = os.path.join(USER_DIR, "image_prompts_cache.json")
|
||||
|
||||
def decode_bytes_safely(raw_bytes):
|
||||
if not isinstance(raw_bytes, bytes):
|
||||
return str(raw_bytes or "")
|
||||
for enc in ["utf-8", "utf-16le", "utf-16", "shift_jis", "gbk", "latin-1"]:
|
||||
try:
|
||||
return raw_bytes.decode(enc)
|
||||
except (UnicodeDecodeError, LookupError):
|
||||
continue
|
||||
return raw_bytes.decode("utf-8", errors="ignore")
|
||||
|
||||
def parse_positive_from_parameters(parameters):
|
||||
if not parameters:
|
||||
return ""
|
||||
lines = parameters.split("\n")
|
||||
if isinstance(parameters, bytes):
|
||||
parameters = decode_bytes_safely(parameters)
|
||||
lines = str(parameters).split("\n")
|
||||
pos_lines = []
|
||||
for line in lines:
|
||||
if line.strip().startswith("Negative prompt:") or line.strip().startswith("Steps:"):
|
||||
@@ -24,31 +49,110 @@ def parse_positive_from_parameters(parameters):
|
||||
pos_lines.append(line)
|
||||
return "\n".join(pos_lines).strip()
|
||||
|
||||
def trace_positive_prompt_from_graph(prompt_data):
|
||||
if not isinstance(prompt_data, dict):
|
||||
return ""
|
||||
|
||||
start_nodes = []
|
||||
for nid, ndata in prompt_data.items():
|
||||
if not isinstance(ndata, dict):
|
||||
continue
|
||||
ctype = ndata.get("class_type", "")
|
||||
inputs = ndata.get("inputs", {})
|
||||
if any(kw in ctype for kw in ["Sampler", "KSampler", "Guider", "CFG"]):
|
||||
for link_key in ["positive", "conditioning", "guider"]:
|
||||
link = inputs.get(link_key)
|
||||
if isinstance(link, list) and len(link) > 0:
|
||||
start_nodes.append(str(link[0]))
|
||||
|
||||
visited = set()
|
||||
found_texts = []
|
||||
|
||||
def walk(nid):
|
||||
if not nid or nid in visited:
|
||||
return
|
||||
visited.add(nid)
|
||||
ndata = prompt_data.get(nid)
|
||||
if not isinstance(ndata, dict):
|
||||
return
|
||||
inputs = ndata.get("inputs", {})
|
||||
|
||||
for tkey in ["text", "prompt", "text_g", "text_l", "text_positive", "positive_prompt", "value", "string"]:
|
||||
val = inputs.get(tkey)
|
||||
if isinstance(val, str) and val.strip():
|
||||
s = val.strip()
|
||||
if s not in found_texts:
|
||||
found_texts.append(s)
|
||||
elif isinstance(val, list) and len(val) > 0:
|
||||
walk(str(val[0]))
|
||||
|
||||
for in_name, in_val in inputs.items():
|
||||
if isinstance(in_val, list) and len(in_val) > 0:
|
||||
low = in_name.lower()
|
||||
if "negative" in low:
|
||||
continue
|
||||
if any(k in low for k in ["conditioning", "positive", "cond", "text", "prompt", "guider"]):
|
||||
walk(str(in_val[0]))
|
||||
|
||||
for sn in start_nodes:
|
||||
walk(sn)
|
||||
|
||||
if found_texts:
|
||||
return "\n".join(found_texts)
|
||||
|
||||
for nid, ndata in prompt_data.items():
|
||||
if not isinstance(ndata, dict):
|
||||
continue
|
||||
ctype = ndata.get("class_type", "")
|
||||
inputs = ndata.get("inputs", {})
|
||||
if any(k in ctype for k in ["CLIPTextEncode", "Text", "Prompt"]):
|
||||
for tkey in ["text", "prompt", "text_g", "text_l"]:
|
||||
val = inputs.get(tkey)
|
||||
if isinstance(val, str) and val.strip() and len(val.strip()) > 1:
|
||||
s = val.strip()
|
||||
if s not in found_texts:
|
||||
found_texts.append(s)
|
||||
return "\n".join(found_texts) if found_texts else ""
|
||||
|
||||
def extract_metadata_from_image(filepath):
|
||||
positive_prompt = ""
|
||||
width, height = 0, 0
|
||||
try:
|
||||
with Image.open(filepath) as img:
|
||||
width, height = img.size
|
||||
if img.format == 'PNG':
|
||||
parameters = img.info.get("parameters", "")
|
||||
info = img.info or {}
|
||||
|
||||
# 1. ComfyUI native format (JSON string in img.info["prompt"])
|
||||
prompt_str = info.get("prompt")
|
||||
if prompt_str:
|
||||
try:
|
||||
prompt_data = json.loads(prompt_str) if isinstance(prompt_str, str) else prompt_str
|
||||
positive_prompt = trace_positive_prompt_from_graph(prompt_data)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 2. A1111 / WebUI parameters
|
||||
if not positive_prompt and "parameters" in info:
|
||||
parameters = info.get("parameters", "")
|
||||
if parameters:
|
||||
positive_prompt = parse_positive_from_parameters(parameters)
|
||||
else:
|
||||
exif_data = img.info.get("exif")
|
||||
|
||||
# 3. EXIF UserComment
|
||||
if not positive_prompt and "exif" in info:
|
||||
exif_data = info.get("exif")
|
||||
if exif_data:
|
||||
exif_dict = piexif.load(exif_data)
|
||||
user_comment = exif_dict.get("Exif", {}).get(piexif.ExifIFD.UserComment)
|
||||
if user_comment:
|
||||
try:
|
||||
try:
|
||||
exif_dict = piexif.load(exif_data)
|
||||
user_comment = exif_dict.get("Exif", {}).get(piexif.ExifIFD.UserComment)
|
||||
if user_comment:
|
||||
parameters = piexif.helper.UserComment.load(user_comment)
|
||||
if isinstance(parameters, bytes):
|
||||
parameters = parameters.decode('utf-8', errors='ignore')
|
||||
positive_prompt = parse_positive_from_parameters(parameters)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Error extracting metadata: {e}")
|
||||
print(f"[SaturnNodes] Error extracting metadata: {e}")
|
||||
return positive_prompt, width, height
|
||||
|
||||
def load_prompts_cache():
|
||||
@@ -62,43 +166,69 @@ def load_prompts_cache():
|
||||
|
||||
def save_prompts_cache(cache):
|
||||
try:
|
||||
with open(CACHE_FILE, "w", encoding="utf-8") as f:
|
||||
temp_file = CACHE_FILE + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(cache, f, indent=4, ensure_ascii=False)
|
||||
os.replace(temp_file, CACHE_FILE)
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Failed to save prompts cache: {e}")
|
||||
print(f"[SaturnNodes] Failed to save prompts cache: {e}")
|
||||
|
||||
server = PromptServer.instance
|
||||
routes = server.routes
|
||||
_PROMPTS_CACHE_LOCK = threading.RLock()
|
||||
IMAGE_THUMBNAILS_DIR = os.path.join(USER_DIR, "cache", "image_thumbs")
|
||||
|
||||
def get_image_metadata_scan_limit():
|
||||
try:
|
||||
limit = int(get_env_setting("IMAGE_LOADER_METADATA_LIMIT", "128"))
|
||||
return max(0, min(2000, limit))
|
||||
except Exception:
|
||||
return 128
|
||||
|
||||
def scan_image_folder(resolved_folder, metadata_limit=None):
|
||||
"""Scan an image folder and read prompt metadata off the event-loop thread."""
|
||||
if metadata_limit is None:
|
||||
metadata_limit = get_image_metadata_scan_limit()
|
||||
|
||||
image_extensions = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
||||
raw_files = []
|
||||
for root, dirs, files in os.walk(resolved_folder):
|
||||
for file in files:
|
||||
ext = os.path.splitext(file)[1]
|
||||
if ext.lower() in image_extensions:
|
||||
full_path = os.path.normpath(os.path.join(root, file))
|
||||
# The folder itself was already validated by the endpoint; just
|
||||
# enforce that walked files stay inside it.
|
||||
if not is_safe_path(full_path, [resolved_folder]):
|
||||
continue
|
||||
try:
|
||||
mtime = os.path.getmtime(full_path)
|
||||
except Exception:
|
||||
mtime = 0
|
||||
rel_path = os.path.relpath(full_path, resolved_folder).replace("\\", "/")
|
||||
raw_files.append((rel_path, full_path, mtime))
|
||||
|
||||
# Sort files by newest modified first. Pre-sort by relative path and use a
|
||||
# stable mtime sort so identical timestamps always produce the same order.
|
||||
raw_files.sort(key=lambda x: x[0])
|
||||
raw_files.sort(key=lambda x: x[2], reverse=True)
|
||||
|
||||
@routes.get("/image_loader/get_images")
|
||||
async def get_images_endpoint(request):
|
||||
folder = request.query.get("folder", "")
|
||||
if not folder or not os.path.exists(folder):
|
||||
return web.json_response({"names": [], "mapping": {}, "prompts": {}})
|
||||
|
||||
image_extensions = {".png", ".jpg", ".jpeg", ".webp", ".PNG", ".JPG", ".JPEG", ".WEBP"}
|
||||
names = []
|
||||
mapping = {}
|
||||
prompts = {}
|
||||
|
||||
cache = load_prompts_cache()
|
||||
cache_modified = False
|
||||
|
||||
for root, dirs, files in os.walk(folder):
|
||||
for file in files:
|
||||
ext = os.path.splitext(file)[1]
|
||||
if ext in image_extensions:
|
||||
full_path = os.path.normpath(os.path.join(root, file))
|
||||
rel_path = os.path.relpath(full_path, folder).replace("\\", "/")
|
||||
names.append(rel_path)
|
||||
mapping[rel_path] = rel_path
|
||||
|
||||
|
||||
with _PROMPTS_CACHE_LOCK:
|
||||
cache = load_prompts_cache()
|
||||
cache_modified = False
|
||||
|
||||
for i, (rel_path, full_path, mtime) in enumerate(raw_files):
|
||||
names.append(rel_path)
|
||||
mapping[rel_path] = rel_path
|
||||
|
||||
if i < metadata_limit:
|
||||
try:
|
||||
mtime = os.path.getmtime(full_path)
|
||||
cached_entry = cache.get(full_path)
|
||||
|
||||
if cached_entry and cached_entry.get("mtime") == mtime:
|
||||
prompt_text = cached_entry.get("prompt", "")
|
||||
cache[full_path] = cache.pop(full_path)
|
||||
else:
|
||||
prompt_text, _, _ = extract_metadata_from_image(full_path)
|
||||
cache[full_path] = {
|
||||
@@ -106,14 +236,80 @@ async def get_images_endpoint(request):
|
||||
"prompt": prompt_text
|
||||
}
|
||||
cache_modified = True
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Error checking cache for {file}: {e}")
|
||||
except Exception:
|
||||
prompt_text = ""
|
||||
|
||||
prompts[rel_path] = prompt_text
|
||||
else:
|
||||
prompts[rel_path] = ""
|
||||
|
||||
if cache_modified:
|
||||
save_prompts_cache(cache)
|
||||
# Prune cache to maximum 256 entries
|
||||
if len(cache) > 256:
|
||||
cache_keys = list(cache.keys())[-256:]
|
||||
cache = {k: cache[k] for k in cache_keys}
|
||||
cache_modified = True
|
||||
|
||||
if cache_modified:
|
||||
save_prompts_cache(cache)
|
||||
|
||||
return names, mapping, prompts
|
||||
|
||||
def get_or_create_image_thumbnail(img_path, mtime, max_size=256):
|
||||
"""Disk-cached WebP thumbnails for the visual image loader."""
|
||||
try:
|
||||
max_size = max(64, min(1024, int(max_size)))
|
||||
cache_key = hashlib.md5(
|
||||
f"{os.path.normcase(os.path.abspath(img_path))}_{mtime}_{max_size}".encode("utf-8")
|
||||
).hexdigest()
|
||||
thumb_path = os.path.join(IMAGE_THUMBNAILS_DIR, f"thumb_{cache_key}.webp")
|
||||
if os.path.exists(thumb_path) and os.path.getsize(thumb_path) > 0:
|
||||
return thumb_path
|
||||
|
||||
with Image.open(img_path) as img:
|
||||
img = ImageOps.exif_transpose(img)
|
||||
if img.width <= max_size and img.height <= max_size:
|
||||
return img_path
|
||||
img = img.convert("RGB")
|
||||
img.thumbnail((max_size, max_size), Image.Resampling.BICUBIC)
|
||||
os.makedirs(IMAGE_THUMBNAILS_DIR, exist_ok=True)
|
||||
# Unique temp names keep concurrent requests for the same image from
|
||||
# overwriting each other's half-written WebP file.
|
||||
temp_path = thumb_path + f".{threading.get_ident()}.tmp"
|
||||
try:
|
||||
img.save(temp_path, format="WEBP", quality=82, method=1)
|
||||
os.replace(temp_path, thumb_path)
|
||||
finally:
|
||||
if os.path.exists(temp_path):
|
||||
try:
|
||||
os.remove(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
return thumb_path
|
||||
except Exception:
|
||||
return img_path
|
||||
|
||||
if hasattr(PromptServer, "instance") and PromptServer.instance is not None:
|
||||
routes = PromptServer.instance.routes
|
||||
else:
|
||||
routes = web.RouteTableDef()
|
||||
print("[SaturnNodes] Notice: PromptServer instance not available during image_loader import; using standalone route table.")
|
||||
|
||||
@routes.get("/saturnnodes/image_loader/get_images")
|
||||
async def get_images_endpoint(request):
|
||||
if not is_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local access only"}, status=403)
|
||||
|
||||
folder = request.query.get("folder", "")
|
||||
if not folder:
|
||||
folder = request.query.get("folder_path", "")
|
||||
|
||||
resolved_folder = sanitize_image_loader_folder(folder)
|
||||
|
||||
if not resolved_folder or not is_safe_path(resolved_folder) or not os.path.exists(resolved_folder):
|
||||
return web.json_response({"names": [], "mapping": {}, "prompts": {}})
|
||||
|
||||
# Folder walks and metadata parsing are blocking disk I/O; keep them off
|
||||
# ComfyUI's asyncio event loop.
|
||||
names, mapping, prompts = await asyncio.to_thread(scan_image_folder, resolved_folder)
|
||||
|
||||
return web.json_response({
|
||||
"names": names,
|
||||
@@ -121,41 +317,73 @@ async def get_images_endpoint(request):
|
||||
"prompts": prompts
|
||||
})
|
||||
|
||||
@routes.get("/image_loader/get_thumbnail")
|
||||
@routes.get("/saturnnodes/image_loader/get_thumbnail")
|
||||
async def get_thumbnail_endpoint(request):
|
||||
if not is_local_request(request):
|
||||
return web.Response(status=403)
|
||||
|
||||
folder = request.query.get("folder", "")
|
||||
if not folder:
|
||||
folder = request.query.get("folder_path", "")
|
||||
resolved_folder = sanitize_image_loader_folder(folder)
|
||||
|
||||
image_rel = request.query.get("image", "")
|
||||
|
||||
full_path = os.path.normpath(os.path.join(folder, image_rel))
|
||||
if not os.path.exists(full_path):
|
||||
if not image_rel:
|
||||
return web.Response(status=400)
|
||||
|
||||
full_path = os.path.normpath(os.path.join(resolved_folder, image_rel))
|
||||
if not is_safe_path(full_path, [resolved_folder]) or not os.path.exists(full_path) or not os.path.isfile(full_path):
|
||||
return web.Response(status=404)
|
||||
|
||||
|
||||
image_extensions = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
||||
if os.path.splitext(full_path)[1].lower() not in image_extensions:
|
||||
return web.Response(status=403)
|
||||
|
||||
try:
|
||||
with Image.open(full_path) as img:
|
||||
img.thumbnail((256, 256))
|
||||
img_byte_arr = io.BytesIO()
|
||||
img.convert("RGB").save(img_byte_arr, format='JPEG', quality=80)
|
||||
return web.Response(body=img_byte_arr.getvalue(), content_type="image/jpeg")
|
||||
thumb_size = int(request.query.get("size", 256))
|
||||
except Exception:
|
||||
thumb_size = 256
|
||||
thumb_size = max(64, min(1024, thumb_size))
|
||||
|
||||
try:
|
||||
mtime = os.path.getmtime(full_path)
|
||||
serve_path = await asyncio.to_thread(get_or_create_image_thumbnail, full_path, mtime, thumb_size)
|
||||
return web.FileResponse(serve_path)
|
||||
except Exception:
|
||||
return web.FileResponse(full_path)
|
||||
|
||||
@routes.get("/image_loader/get_full_image")
|
||||
@routes.get("/saturnnodes/image_loader/get_full_image")
|
||||
async def get_full_image_endpoint(request):
|
||||
if not is_local_request(request):
|
||||
return web.Response(status=403)
|
||||
|
||||
folder = request.query.get("folder", "")
|
||||
if not folder:
|
||||
folder = request.query.get("folder_path", "")
|
||||
resolved_folder = sanitize_image_loader_folder(folder)
|
||||
|
||||
image_rel = request.query.get("image", "")
|
||||
|
||||
full_path = os.path.normpath(os.path.join(folder, image_rel))
|
||||
if not os.path.exists(full_path):
|
||||
if not image_rel:
|
||||
return web.Response(status=400)
|
||||
|
||||
full_path = os.path.normpath(os.path.join(resolved_folder, image_rel))
|
||||
if not is_safe_path(full_path, [resolved_folder]) or not os.path.exists(full_path) or not os.path.isfile(full_path):
|
||||
return web.Response(status=404)
|
||||
|
||||
image_extensions = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
||||
if os.path.splitext(full_path)[1].lower() not in image_extensions:
|
||||
return web.Response(status=403)
|
||||
|
||||
return web.FileResponse(full_path)
|
||||
|
||||
class ImageLoaderVisualPrettyV2:
|
||||
class VisualImageLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"folder_path": ("STRING", {"default": ""}),
|
||||
"display_mode": (["Scrollable", "Show All"], {"default": "Scrollable"}),
|
||||
"folder": ("STRING", {"default": ""}),
|
||||
"display_mode": (["Scrollable", "Show All"], {"default": "Scrollable", "advanced": True}),
|
||||
"sort_images_by": (["Name (A-Z)", "Name (Z-A)", "Date Modified (Newest First)", "Date Modified (Oldest First)"], {"default": "Date Modified (Newest First)", "advanced": True}),
|
||||
},
|
||||
"hidden": {
|
||||
"_selected_image": ("STRING", {"default": ""}),
|
||||
@@ -172,25 +400,29 @@ class ImageLoaderVisualPrettyV2:
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return kwargs.get("_selected_image", "")
|
||||
|
||||
def load_image(self, folder_path, display_mode="Scrollable", _selected_image=""):
|
||||
if not folder_path or not _selected_image:
|
||||
empty_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
return (empty_image, "", 0, 0)
|
||||
|
||||
full_path = os.path.normpath(os.path.join(folder_path, _selected_image))
|
||||
if not os.path.exists(full_path):
|
||||
empty_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
return (empty_image, "", 0, 0)
|
||||
|
||||
i = Image.open(full_path)
|
||||
i = ImageOps.exif_transpose(i)
|
||||
def load_image(self, folder="", display_mode="Scrollable", sort_images_by="Date Modified (Newest First)", folder_path=None, _selected_image="", **kwargs):
|
||||
input_dir = folder if folder else (folder_path or "")
|
||||
resolved_folder = sanitize_image_loader_folder(input_dir)
|
||||
|
||||
positive_prompt, width, height = extract_metadata_from_image(full_path)
|
||||
if width == 0 or height == 0:
|
||||
width, height = i.size
|
||||
if not resolved_folder or not _selected_image:
|
||||
empty_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
return (empty_image, "", 0, 0)
|
||||
|
||||
image = i.convert("RGB")
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)[None,]
|
||||
full_path = os.path.normpath(os.path.join(resolved_folder, _selected_image))
|
||||
if not is_safe_path(full_path) or not os.path.exists(full_path):
|
||||
empty_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
return (empty_image, "", 0, 0)
|
||||
|
||||
with Image.open(full_path) as i:
|
||||
i = ImageOps.exif_transpose(i)
|
||||
|
||||
positive_prompt, width, height = extract_metadata_from_image(full_path)
|
||||
if width == 0 or height == 0:
|
||||
width, height = i.size
|
||||
|
||||
image = i.convert("RGB")
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)[None,]
|
||||
|
||||
return (image, positive_prompt, width, height)
|
||||
|
||||
|
||||
+50
-27
@@ -3,8 +3,9 @@ import torch
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps
|
||||
import folder_paths
|
||||
from .utils import sanitize_folder_path
|
||||
|
||||
LOADER_CATEGORY = "🍃 FlowControl/Loaders"
|
||||
LOADER_CATEGORY = "🪐 SaturnNodes/Loaders"
|
||||
|
||||
class LoadRecentOutputs:
|
||||
@classmethod
|
||||
@@ -20,55 +21,77 @@ class LoadRecentOutputs:
|
||||
"step": 1,
|
||||
"control_after_generate": "increment"
|
||||
}),
|
||||
"order": (["Newest First (0 = latest)", "Chronological (Oldest First)"], {
|
||||
"default": "Newest First (0 = latest)",
|
||||
"advanced": True
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "load_single_image"
|
||||
CATEGORY = LOADER_CATEGORY
|
||||
DESCRIPTION = "Cycles through recently generated output images.\n\nExplanation of Index:\n- 'amount' grabs the X newest images from the folder.\n- 'index' selects which of those X images to output.\n- For example, if amount=64, setting index=0 to 63 in a batch will load each of the 64 newest images one by one.\n- If index exceeds amount, it loops back around (e.g. index 64 = index 0)."
|
||||
DESCRIPTION = "Cycles through recently generated output images.\n\nExplanation of Index:\n- 'amount' grabs the X newest images from the folder.\n- 'index' selects which of those X images to output (0 = newest image).\n- If index exceeds amount, it loops back around."
|
||||
|
||||
def load_single_image(self, output_folder, amount, index):
|
||||
if output_folder == "output":
|
||||
def load_single_image(self, output_folder, amount, index, order="Newest First (0 = latest)", **kwargs):
|
||||
clean_input = (output_folder or "").strip()
|
||||
if clean_input == "output" or not clean_input:
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
else:
|
||||
output_dir = output_folder
|
||||
output_dir = sanitize_folder_path(clean_input, default_dir=folder_paths.get_output_directory())
|
||||
|
||||
if not os.path.exists(output_dir):
|
||||
dummy = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
dummy = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
return (dummy,)
|
||||
|
||||
valid_extensions = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
|
||||
file_list = []
|
||||
for root, _, files in os.walk(output_dir):
|
||||
for file in files:
|
||||
if file.lower().endswith(valid_extensions):
|
||||
path = os.path.join(root, file)
|
||||
try:
|
||||
mtime = os.path.getmtime(path)
|
||||
file_list.append((path, mtime))
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
if not file_list:
|
||||
dummy = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
try:
|
||||
for root, dirs, files in os.walk(output_dir):
|
||||
for file in files:
|
||||
if file.lower().endswith(valid_extensions):
|
||||
path = os.path.join(root, file)
|
||||
try:
|
||||
mtime = os.path.getmtime(path)
|
||||
file_list.append((path, mtime))
|
||||
except OSError:
|
||||
continue
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error scanning output folder '{output_dir}': {e}")
|
||||
dummy = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
return (dummy,)
|
||||
|
||||
file_list.sort(key=lambda x: x[1], reverse=True)
|
||||
if not file_list:
|
||||
dummy = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
return (dummy,)
|
||||
|
||||
try:
|
||||
amount = max(1, int(amount))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
amount = 1
|
||||
try:
|
||||
index = max(0, int(index))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
index = 0
|
||||
|
||||
# Newest first, with the path as a deterministic tie-breaker for files
|
||||
# copied/extracted with identical timestamps.
|
||||
file_list.sort(key=lambda x: (x[1], x[0]), reverse=True)
|
||||
recent_files = file_list[:amount]
|
||||
recent_files.reverse()
|
||||
if "oldest" in str(order).lower() or "chronological" in str(order).lower():
|
||||
recent_files.reverse()
|
||||
|
||||
actual_index = index % len(recent_files)
|
||||
target_file, _ = recent_files[actual_index]
|
||||
|
||||
try:
|
||||
img = Image.open(target_file)
|
||||
img = ImageOps.exif_transpose(img)
|
||||
image = img.convert("RGB")
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image).unsqueeze(0)
|
||||
with Image.open(target_file) as img:
|
||||
img = ImageOps.exif_transpose(img)
|
||||
image = img.convert("RGB")
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image).unsqueeze(0)
|
||||
except Exception as e:
|
||||
print(f"[FlowControl] Error loading recent output image {target_file}: {e}")
|
||||
image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
print(f"[SaturnNodes] Error loading recent output image {target_file}: {e}")
|
||||
image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||||
|
||||
return (image,)
|
||||
|
||||
@@ -0,0 +1,449 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import shlex
|
||||
import subprocess
|
||||
import threading
|
||||
from aiohttp import web
|
||||
import folder_paths
|
||||
import comfy.model_management
|
||||
from .utils import get_env_setting, is_local_request, is_authenticated_local_request
|
||||
|
||||
AUTOMATION_CATEGORY = "🪐 SaturnNodes/Automation"
|
||||
|
||||
def _safe_print(msg):
|
||||
try:
|
||||
print(msg)
|
||||
except UnicodeEncodeError:
|
||||
print(msg.encode("ascii", "replace").decode("ascii"))
|
||||
|
||||
def is_local_file_execution_enabled():
|
||||
"""Checks whether local file execution is explicitly enabled in SaturnNodes settings (default False)."""
|
||||
val = get_env_setting("ENABLE_LOCAL_FILE_EXECUTION", "false").lower()
|
||||
return val in ["true", "1", "yes"]
|
||||
|
||||
def get_allowed_scripts_directory():
|
||||
"""
|
||||
Returns the strict, non-bypassable scripts directory: ComfyUI/scripts/
|
||||
"""
|
||||
base = folder_paths.base_path if hasattr(folder_paths, "base_path") and folder_paths.base_path else os.getcwd()
|
||||
scripts_dir = os.path.abspath(os.path.realpath(os.path.join(base, "scripts")))
|
||||
try:
|
||||
os.makedirs(scripts_dir, exist_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
return scripts_dir
|
||||
|
||||
def is_script_path_permitted(target_path):
|
||||
"""
|
||||
Confines target scripts strictly to ComfyUI/scripts/ using canonical realpath and commonpath.
|
||||
No setting or configuration can bypass this confinement.
|
||||
"""
|
||||
if not target_path:
|
||||
return False, "Target path is empty."
|
||||
|
||||
scripts_dir = os.path.abspath(os.path.realpath(get_allowed_scripts_directory()))
|
||||
try:
|
||||
abs_target = os.path.abspath(os.path.realpath(target_path))
|
||||
norm_scripts = os.path.normcase(scripts_dir)
|
||||
norm_target = os.path.normcase(abs_target)
|
||||
if os.path.commonpath([norm_scripts, norm_target]) == norm_scripts:
|
||||
if not os.path.exists(abs_target):
|
||||
return False, f"Script does not exist: '{os.path.basename(target_path)}'"
|
||||
if not os.path.isfile(abs_target):
|
||||
return False, f"Target is not a regular file: '{os.path.basename(target_path)}'"
|
||||
return True, "Path verified within ComfyUI/scripts/."
|
||||
return False, f"Security Restriction: Script '{os.path.basename(target_path)}' must reside within ComfyUI/scripts/."
|
||||
except Exception as e:
|
||||
return False, f"Path validation error: {e}"
|
||||
|
||||
def resolve_target_path(file_path_input):
|
||||
"""
|
||||
Resolves and verifies a script path strictly within ComfyUI/scripts/.
|
||||
Rejects absolute paths, path traversal ('..'), and any outside references.
|
||||
"""
|
||||
clean = str(file_path_input or "").strip()
|
||||
if not clean:
|
||||
return None, "Script file path is empty."
|
||||
|
||||
# 1. Strictly reject absolute paths (Windows drive letters C:\, UNC \\server, or POSIX /root)
|
||||
has_drive_letter = len(clean) >= 2 and clean[0].isalpha() and clean[1] == ":"
|
||||
if os.path.isabs(clean) or clean.startswith(("\\\\", "//", "/")) or has_drive_letter:
|
||||
return None, f"Security Restriction: Absolute paths are forbidden ('{clean}'). Place your script in ComfyUI/scripts/ and specify a relative name."
|
||||
|
||||
# 2. Strictly reject directory traversal
|
||||
norm = os.path.normpath(clean)
|
||||
parts = norm.replace("\\", "/").split("/")
|
||||
if ".." in parts or norm.startswith(".."):
|
||||
return None, f"Security Restriction: Directory traversal ('..') is forbidden ('{clean}')."
|
||||
|
||||
# 3. If operator wrote "scripts/run.bat" or "scripts\run.bat", normalize to relative subpath
|
||||
if parts and parts[0].lower() == "scripts":
|
||||
clean_sub = os.path.join(*parts[1:]) if len(parts) > 1 else ""
|
||||
else:
|
||||
clean_sub = norm
|
||||
|
||||
if not clean_sub:
|
||||
return None, "Invalid script path."
|
||||
|
||||
scripts_base = os.path.abspath(os.path.realpath(get_allowed_scripts_directory()))
|
||||
full_target = os.path.abspath(os.path.realpath(os.path.join(scripts_base, clean_sub)))
|
||||
|
||||
# 4. Canonical commonpath confinement
|
||||
try:
|
||||
norm_base = os.path.normcase(scripts_base)
|
||||
norm_target = os.path.normcase(full_target)
|
||||
if os.path.commonpath([norm_base, norm_target]) != norm_base:
|
||||
return None, f"Security Restriction: Script path escapes ComfyUI/scripts/ ('{clean}')."
|
||||
except Exception as e:
|
||||
return None, f"Security Restriction: Path validation error: {e}"
|
||||
|
||||
if not os.path.exists(full_target):
|
||||
return None, f"Script not found in ComfyUI/scripts/: '{clean_sub}'"
|
||||
|
||||
if not os.path.isfile(full_target):
|
||||
return None, f"Target path is not a file: '{clean_sub}'"
|
||||
|
||||
return full_target, "OK"
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# In-Memory Operator Authorization Registry
|
||||
# =========================================================================
|
||||
# Does NOT travel in workflow JSON or prompt requests.
|
||||
# Stored solely in memory and consumed upon execution.
|
||||
_AUTHORIZED_RUNS = {}
|
||||
_AUTH_LOCK = threading.Lock()
|
||||
|
||||
def authorize_node_run(node_id, file_path="", ttl_seconds=300):
|
||||
with _AUTH_LOCK:
|
||||
now = time.time()
|
||||
expired = [k for k, v in _AUTHORIZED_RUNS.items() if v.get("expires_at", 0) < now]
|
||||
for k in expired:
|
||||
_AUTHORIZED_RUNS.pop(k, None)
|
||||
|
||||
filename = os.path.basename(str(file_path).strip()) if file_path else ""
|
||||
_AUTHORIZED_RUNS[str(node_id)] = {
|
||||
"expires_at": now + ttl_seconds,
|
||||
"filename": filename
|
||||
}
|
||||
|
||||
def consume_node_authorization(node_id, resolved_path):
|
||||
if node_id is None:
|
||||
return False, "Missing node identifier. Workflow execution cannot be authorized."
|
||||
with _AUTH_LOCK:
|
||||
now = time.time()
|
||||
auth = _AUTHORIZED_RUNS.pop(str(node_id), None)
|
||||
if not auth:
|
||||
return False, "Node has not been authorized. Operator must click 'Authorize Run' on this node in the canvas before queueing."
|
||||
if auth.get("expires_at", 0) < now:
|
||||
return False, "Authorization has expired (5-minute timeout). Please click 'Authorize Run' again."
|
||||
target_base = os.path.basename(resolved_path) if resolved_path else ""
|
||||
expected_base = auth.get("filename")
|
||||
if expected_base and target_base and target_base.lower() != expected_base.lower():
|
||||
return False, f"Authorized script '{expected_base}' does not match configured script '{target_base}'. Please re-authorize."
|
||||
return True, "Authorized"
|
||||
|
||||
def is_node_authorized(node_id):
|
||||
if not node_id:
|
||||
return False, 0
|
||||
with _AUTH_LOCK:
|
||||
now = time.time()
|
||||
auth = _AUTHORIZED_RUNS.get(str(node_id))
|
||||
if auth and auth.get("expires_at", 0) >= now:
|
||||
return True, int(auth.get("expires_at", 0) - now)
|
||||
return False, 0
|
||||
|
||||
def clear_all_authorizations():
|
||||
with _AUTH_LOCK:
|
||||
_AUTHORIZED_RUNS.clear()
|
||||
|
||||
def setup_local_runner_routes(server):
|
||||
routes = server.routes
|
||||
|
||||
@routes.post("/saturnnodes/local_runner/authorize")
|
||||
async def authorize_endpoint(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
if not is_local_file_execution_enabled():
|
||||
return web.json_response({
|
||||
"error": "Local file execution is disabled. Enable 'Allow Local File Execution' in SaturnNodes settings first.",
|
||||
"enabled": False
|
||||
}, status=403)
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception:
|
||||
data = {}
|
||||
node_id = data.get("node_id")
|
||||
file_path = data.get("file_path", "")
|
||||
if node_id is None:
|
||||
return web.json_response({"error": "Missing node_id"}, status=400)
|
||||
authorize_node_run(node_id, file_path, ttl_seconds=300)
|
||||
return web.json_response({
|
||||
"success": True,
|
||||
"node_id": str(node_id),
|
||||
"expires_in": 300
|
||||
})
|
||||
|
||||
@routes.get("/saturnnodes/local_runner/status")
|
||||
async def status_endpoint(request):
|
||||
if not is_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local access only"}, status=403)
|
||||
node_id = request.query.get("node_id")
|
||||
is_auth, remaining = is_node_authorized(node_id)
|
||||
return web.json_response({
|
||||
"enabled": is_local_file_execution_enabled(),
|
||||
"authorized": is_auth,
|
||||
"remaining_seconds": remaining
|
||||
})
|
||||
|
||||
|
||||
class AnyType(str):
|
||||
"""Wildcard type for input and output passthrough wires."""
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
any_type = AnyType("*")
|
||||
|
||||
|
||||
class RunLocalFileNode:
|
||||
"""
|
||||
Safely executes a local file (.exe, .bat, .cmd, .ps1 on Windows; .sh or binary on Linux/macOS)
|
||||
with custom parameters and configurable working directory.
|
||||
Includes strict path confinement and in-memory operator authorization.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"file_path": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "Relative script path inside ComfyUI/scripts/ (e.g. run.bat or convert.sh)",
|
||||
"tooltip": "Script or binary to execute. MUST be located inside 'ComfyUI/scripts/'. Subfolders (e.g. tools/process.bat) are permitted. Absolute paths and directory traversal ('..') are strictly blocked."
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"parameters": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"placeholder": "Command-line arguments (e.g. --input file.png --quality 95)",
|
||||
"tooltip": "Optional command-line arguments passed to the script. Parsed safely without shell expansion."
|
||||
}),
|
||||
"working_directory": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "Working directory inside ComfyUI/scripts/. Leave empty to use script directory.",
|
||||
"tooltip": "Execution working directory. MUST be located inside 'ComfyUI/scripts/'. Leave empty to default to the script's own folder."
|
||||
}),
|
||||
"run_mode": (
|
||||
["Synchronous (Wait for Output)", "Asynchronous (Background)"],
|
||||
{"default": "Synchronous (Wait for Output)"}
|
||||
),
|
||||
"timeout": ("INT", {
|
||||
"default": 60,
|
||||
"min": 0,
|
||||
"max": 3600,
|
||||
"step": 5,
|
||||
"tooltip": "Max execution time in seconds (0 = unlimited). Ignored in asynchronous mode."
|
||||
}),
|
||||
"trigger": (any_type, {"tooltip": "Optional passthrough wire to sequence execution in your workflow."}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO"
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "INT", "BOOLEAN", any_type)
|
||||
RETURN_NAMES = ("stdout", "stderr", "exit_code", "success", "passthrough")
|
||||
FUNCTION = "run_local_file"
|
||||
CATEGORY = AUTOMATION_CATEGORY
|
||||
DESCRIPTION = (
|
||||
"Executes a local script or executable (.bat, .ps1, .exe on Windows; .sh or binary on Linux/macOS) "
|
||||
"with parameter arguments and working directory control.\n\n"
|
||||
"Security Safeguards:\n"
|
||||
"- Strictly confined to ComfyUI/scripts/ (absolute paths and directory traversal are forbidden).\n"
|
||||
"- Requires explicit in-memory operator authorization via the '⚡ Authorize Run' button.\n"
|
||||
"- Consent cannot travel in workflow JSON or image metadata.\n"
|
||||
"- Execution is single-use: each authorization expires upon run or after 5 minutes.\n"
|
||||
"- Never uses shell=True to prevent command chaining injection.\n"
|
||||
"- Disabled globally by default for operator security."
|
||||
)
|
||||
|
||||
def resolve_target_path(self, file_path_input):
|
||||
return resolve_target_path(file_path_input)[0]
|
||||
|
||||
def build_command_args(self, target_path, params_str):
|
||||
"""Builds a safe, non-shell command list suitable for subprocess.Popen(shell=False)."""
|
||||
parsed_args = []
|
||||
if params_str and params_str.strip():
|
||||
try:
|
||||
parsed_args = shlex.split(params_str.strip(), posix=(sys.platform != "win32"))
|
||||
except Exception:
|
||||
parsed_args = params_str.strip().split()
|
||||
|
||||
norm_path = os.path.normpath(target_path)
|
||||
lower_path = norm_path.lower()
|
||||
|
||||
if sys.platform.startswith("win"):
|
||||
if lower_path.endswith((".bat", ".cmd")):
|
||||
comspec = os.environ.get("COMSPEC", "cmd.exe")
|
||||
return [comspec, "/c", norm_path, *parsed_args]
|
||||
elif lower_path.endswith(".ps1"):
|
||||
# Security Note: -ExecutionPolicy Bypass is required on Windows to permit execution
|
||||
# of local user scripts (.ps1) which are blocked by default machine policies.
|
||||
# Security is strictly guaranteed by:
|
||||
# 1. Path confinement: target file MUST reside exclusively inside ComfyUI/scripts/
|
||||
# 2. Live in-memory operator authorization required on canvas prior to run
|
||||
# 3. -NoProfile flag ensuring no untrusted global PowerShell profiles are executed
|
||||
return ["powershell.exe", "-NoProfile", "-ExecutionPolicy", "Bypass", "-File", norm_path, *parsed_args]
|
||||
else:
|
||||
return [norm_path, *parsed_args]
|
||||
else:
|
||||
if lower_path.endswith(".sh"):
|
||||
return ["/bin/bash", norm_path, *parsed_args]
|
||||
else:
|
||||
return [norm_path, *parsed_args]
|
||||
|
||||
def run_local_file(
|
||||
self,
|
||||
file_path,
|
||||
parameters="",
|
||||
working_directory="",
|
||||
run_mode="Synchronous (Wait for Output)",
|
||||
timeout=60,
|
||||
trigger=None,
|
||||
unique_id=None,
|
||||
**kwargs
|
||||
):
|
||||
passthrough_out = trigger if trigger is not None else ""
|
||||
|
||||
# Check 1: Global Security Setting (Defaults to Disabled)
|
||||
if not is_local_file_execution_enabled():
|
||||
err = (
|
||||
"Execution Blocked: Local file execution is disabled by default for security. "
|
||||
"To run local files, enable 'Allow Local File Execution' in ComfyUI Settings -> SaturnNodes -> Security."
|
||||
)
|
||||
_safe_print(f"[SaturnNodes LocalRunner] {err}")
|
||||
return ("", err, -1, False, passthrough_out)
|
||||
|
||||
# Check 2: Path Resolution & Strict Confinement to ComfyUI/scripts/
|
||||
resolved_path, path_msg = resolve_target_path(file_path)
|
||||
if not resolved_path:
|
||||
err = f"Execution Error: {path_msg}"
|
||||
_safe_print(f"[SaturnNodes LocalRunner] {err}")
|
||||
return ("", err, -1, False, passthrough_out)
|
||||
|
||||
# Check 3: Mandatory In-Memory Operator Authorization Gate
|
||||
node_id = unique_id if unique_id is not None else kwargs.get("node_id")
|
||||
authorized, auth_msg = consume_node_authorization(node_id, resolved_path)
|
||||
if not authorized:
|
||||
err = f"Execution Blocked: {auth_msg}"
|
||||
_safe_print(f"[SaturnNodes LocalRunner] {err}")
|
||||
return ("", err, -1, False, passthrough_out)
|
||||
|
||||
# Determine working directory (confined to ComfyUI/scripts/ or script directory)
|
||||
scripts_base = get_allowed_scripts_directory()
|
||||
cwd = os.path.dirname(resolved_path)
|
||||
if working_directory and working_directory.strip():
|
||||
cand_cwd = os.path.abspath(os.path.realpath(os.path.join(scripts_base, working_directory.strip())))
|
||||
try:
|
||||
if os.path.commonpath([scripts_base, cand_cwd]) == scripts_base and os.path.isdir(cand_cwd):
|
||||
cwd = cand_cwd
|
||||
else:
|
||||
_safe_print(f"[SaturnNodes LocalRunner] Warning: working_directory '{working_directory}' outside scripts folder. Defaulting to script folder.")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
cmd_list = self.build_command_args(resolved_path, parameters)
|
||||
_safe_print(f"[SaturnNodes LocalRunner] Executing '{os.path.basename(resolved_path)}' (mode: {run_mode}, cwd: '{cwd}')...")
|
||||
|
||||
# Mode: Asynchronous (Background Detached)
|
||||
if "Asynchronous" in run_mode:
|
||||
try:
|
||||
creationflags = 0
|
||||
if sys.platform.startswith("win"):
|
||||
creationflags = subprocess.CREATE_NEW_CONSOLE if hasattr(subprocess, "CREATE_NEW_CONSOLE") else 0
|
||||
subprocess.Popen(cmd_list, cwd=cwd, creationflags=creationflags)
|
||||
else:
|
||||
subprocess.Popen(cmd_list, cwd=cwd, start_new_session=True)
|
||||
|
||||
msg = f"Started in background: '{os.path.basename(resolved_path)}' (PID assigned)"
|
||||
_safe_print(f"[SaturnNodes LocalRunner] {msg}")
|
||||
return (msg, "", 0, True, passthrough_out)
|
||||
except Exception as e:
|
||||
err_msg = f"Failed to launch background process: {str(e)}"
|
||||
_safe_print(f"[SaturnNodes LocalRunner] Error: {err_msg}")
|
||||
return ("", err_msg, -1, False, passthrough_out)
|
||||
|
||||
# Mode: Synchronous (Wait for Process with Output & Interrupt Polling)
|
||||
try:
|
||||
creationflags = 0
|
||||
if sys.platform.startswith("win"):
|
||||
creationflags = subprocess.CREATE_NO_WINDOW if hasattr(subprocess, "CREATE_NO_WINDOW") else 0
|
||||
|
||||
proc = subprocess.Popen(
|
||||
cmd_list,
|
||||
cwd=cwd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
creationflags=creationflags
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
def _kill_tree():
|
||||
try:
|
||||
if sys.platform.startswith("win"):
|
||||
subprocess.run(["taskkill", "/F", "/T", "/PID", str(proc.pid)], capture_output=True)
|
||||
else:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
proc.communicate(timeout=1)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
effective_timeout = timeout if timeout > 0 else None
|
||||
stdout_data = ""
|
||||
stderr_data = ""
|
||||
|
||||
# Continuously drain pipes to prevent OS buffer deadlocks while honoring cancellation
|
||||
while True:
|
||||
try:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
except Exception:
|
||||
_kill_tree()
|
||||
_safe_print(f"[SaturnNodes LocalRunner] Execution of '{os.path.basename(resolved_path)}' interrupted by user.")
|
||||
raise
|
||||
|
||||
if effective_timeout and (time.time() - start_time) > effective_timeout:
|
||||
_kill_tree()
|
||||
err = f"Execution timed out after {timeout} seconds."
|
||||
_safe_print(f"[SaturnNodes LocalRunner] {err}")
|
||||
return ("", err, -1, False, passthrough_out)
|
||||
|
||||
try:
|
||||
stdout_data, stderr_data = proc.communicate(timeout=0.1)
|
||||
break
|
||||
except subprocess.TimeoutExpired:
|
||||
continue
|
||||
|
||||
exit_code = proc.returncode
|
||||
success = (exit_code == 0)
|
||||
|
||||
_safe_print(f"[SaturnNodes LocalRunner] Completed with exit code {exit_code} (success: {success}).")
|
||||
return (stdout_data or "", stderr_data or "", exit_code, success, passthrough_out)
|
||||
|
||||
except Exception as e:
|
||||
err_msg = f"Process execution error: {str(e)}"
|
||||
_safe_print(f"[SaturnNodes LocalRunner] {err_msg}")
|
||||
return ("", err_msg, -1, False, passthrough_out)
|
||||
@@ -0,0 +1,188 @@
|
||||
import os
|
||||
import re
|
||||
import folder_paths
|
||||
import comfy.sd
|
||||
from comfy.utils import load_torch_file
|
||||
from .lora_loader import (
|
||||
get_filtered_loras_mapping,
|
||||
increment_lora_usage,
|
||||
LORA_CATEGORY,
|
||||
LORA_OUTPUT_FORMAT_CHOICES
|
||||
)
|
||||
from .utils import parse_pretty_name, format_lora_output_name
|
||||
|
||||
class TextLoraFinder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"clip": ("CLIP",),
|
||||
"folder": ("STRING", {"default": ""}),
|
||||
"search_for": ([
|
||||
"Parsed Name",
|
||||
"Filename",
|
||||
"Filename without extension",
|
||||
"Custom Regex"
|
||||
], {"default": "Parsed Name"}),
|
||||
"custom_regex": ("STRING", {"default": "", "advanced": True}),
|
||||
"search_mode": (["First match (Front)", "Last match (Back)"], {"default": "First match (Front)"}),
|
||||
"find_amount": ("INT", {"default": 1, "min": 1, "max": 20, "step": 1}),
|
||||
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
||||
"strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
||||
"output_format": (LORA_OUTPUT_FORMAT_CHOICES, {"default": "Filename", "advanced": True}),
|
||||
},
|
||||
"optional": {
|
||||
"text": ("STRING", {"forceInput": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "STRING")
|
||||
RETURN_NAMES = ("MODEL", "CLIP", "loras")
|
||||
FUNCTION = "find_and_load_loras"
|
||||
CATEGORY = LORA_CATEGORY
|
||||
DESCRIPTION = "Scans input text for LoRA names, automatically loads matched LoRAs into Model/CLIP, and outputs the formatted list."
|
||||
|
||||
def find_and_load_loras(
|
||||
self,
|
||||
model,
|
||||
clip,
|
||||
folder="",
|
||||
search_for="Parsed Name",
|
||||
custom_regex="",
|
||||
search_mode="First match (Front)",
|
||||
find_amount=1,
|
||||
strength_model=1.0,
|
||||
strength_clip=1.0,
|
||||
output_format="Filename",
|
||||
text=None,
|
||||
start_from=None,
|
||||
prompt=None,
|
||||
**kwargs
|
||||
):
|
||||
active_mode = search_mode
|
||||
if start_from is not None and search_mode == "First match (Front)":
|
||||
if start_from == "Back":
|
||||
active_mode = "Last match (Back)"
|
||||
elif start_from == "Front":
|
||||
active_mode = "First match (Front)"
|
||||
|
||||
input_text = text if text is not None else prompt
|
||||
text_str = str(input_text) if input_text is not None else ""
|
||||
if not text_str:
|
||||
return (model, clip, "")
|
||||
|
||||
mapping = get_filtered_loras_mapping(folder, pretty=True)
|
||||
if not mapping:
|
||||
return (model, clip, "")
|
||||
|
||||
matches = []
|
||||
seen_paths = set()
|
||||
|
||||
if search_for == "Custom Regex":
|
||||
if custom_regex and custom_regex.strip():
|
||||
# Invert the name lookup: build it once instead of scanning the
|
||||
# whole LoRA mapping for every regex match in the prompt.
|
||||
lookup = {}
|
||||
for display_name, sys_path in mapping.items():
|
||||
if display_name in ["[ NONE ]", "[ RANDOM ]"] or not sys_path or sys_path == "[ NONE ]":
|
||||
continue
|
||||
pretty_val = re.sub(r'\s+V\d+(\.\d+)?$', '', parse_pretty_name(sys_path), flags=re.IGNORECASE).strip()
|
||||
fname = os.path.basename(sys_path)
|
||||
fname_noext = os.path.splitext(fname)[0]
|
||||
for key in (pretty_val, fname, fname_noext, display_name):
|
||||
if key:
|
||||
lookup.setdefault(key.lower(), []).append((display_name, sys_path))
|
||||
try:
|
||||
pattern = re.compile(custom_regex, re.IGNORECASE)
|
||||
for m in pattern.finditer(text_str):
|
||||
matched_text = m.group(0).strip().lower()
|
||||
for display_name, sys_path in lookup.get(matched_text, []):
|
||||
if sys_path in seen_paths:
|
||||
continue
|
||||
matches.append((m.start(), display_name, sys_path))
|
||||
seen_paths.add(sys_path)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Invalid custom regex pattern '{custom_regex}': {e}")
|
||||
else:
|
||||
# Build one name lookup and scan the prompt once (combined regex),
|
||||
# instead of running one regex scan per installed LoRA.
|
||||
lookup = {}
|
||||
for display_name, sys_path in mapping.items():
|
||||
if display_name in ["[ NONE ]", "[ RANDOM ]"] or not sys_path or sys_path == "[ NONE ]":
|
||||
continue
|
||||
|
||||
if search_for in ["Pretty Name", "Parsed Name"]:
|
||||
target_name = re.sub(r'\s+V\d+(\.\d+)?$', '', parse_pretty_name(sys_path), flags=re.IGNORECASE).strip()
|
||||
elif search_for == "Filename":
|
||||
target_name = os.path.basename(sys_path)
|
||||
else: # Filename without extension
|
||||
target_name = os.path.splitext(os.path.basename(sys_path))[0]
|
||||
|
||||
if not target_name:
|
||||
continue
|
||||
lookup.setdefault(target_name.lower(), []).append((display_name, sys_path))
|
||||
|
||||
if lookup:
|
||||
# Longest-first so "Style XL" wins over "Style" at the same position.
|
||||
ordered_names = sorted(lookup.keys(), key=len, reverse=True)
|
||||
combined_pattern = r'(?<!\w)(?:' + "|".join(re.escape(n) for n in ordered_names) + r')(?!\w)'
|
||||
try:
|
||||
combined_regex = re.compile(combined_pattern, flags=re.IGNORECASE)
|
||||
for m in combined_regex.finditer(text_str):
|
||||
for display_name, sys_path in lookup.get(m.group(0).lower(), []):
|
||||
if sys_path in seen_paths:
|
||||
continue
|
||||
matches.append((m.start(), display_name, sys_path))
|
||||
seen_paths.add(sys_path)
|
||||
except re.error:
|
||||
# Extremely large catalogs can exceed the regex compiler's
|
||||
# limits; fall back to the original per-name scan.
|
||||
for target_name, entries in lookup.items():
|
||||
pattern = r'(?<!\w)' + re.escape(target_name) + r'(?!\w)'
|
||||
for m in re.finditer(pattern, text_str, flags=re.IGNORECASE):
|
||||
for display_name, sys_path in entries:
|
||||
if sys_path in seen_paths:
|
||||
continue
|
||||
matches.append((m.start(), display_name, sys_path))
|
||||
seen_paths.add(sys_path)
|
||||
|
||||
if not matches:
|
||||
return (model, clip, "")
|
||||
|
||||
matches.sort(key=lambda x: x[0])
|
||||
|
||||
num_to_take = max(1, int(find_amount))
|
||||
if active_mode in ["Back", "Last match (Back)"]:
|
||||
selected_matches = matches[-num_to_take:]
|
||||
else:
|
||||
selected_matches = matches[:num_to_take]
|
||||
|
||||
current_model = model
|
||||
current_clip = clip
|
||||
formatted_names = []
|
||||
|
||||
for start_idx, disp_name, sys_path in selected_matches:
|
||||
lora_path = folder_paths.get_full_path("loras", sys_path)
|
||||
if not lora_path or not os.path.exists(lora_path):
|
||||
continue
|
||||
|
||||
if strength_model != 0 or strength_clip != 0:
|
||||
try:
|
||||
lora_weights = load_torch_file(lora_path, safe_load=True)
|
||||
current_model, current_clip = comfy.sd.load_lora_for_models(
|
||||
current_model, current_clip, lora_weights, strength_model, strength_clip
|
||||
)
|
||||
increment_lora_usage(sys_path)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error loading LoRA '{sys_path}': {e}")
|
||||
continue
|
||||
else:
|
||||
increment_lora_usage(sys_path)
|
||||
|
||||
out_name = format_lora_output_name(sys_path, disp_name, output_format=output_format, custom_regex=custom_regex)
|
||||
if out_name:
|
||||
formatted_names.append(out_name)
|
||||
|
||||
return (current_model, current_clip, ", ".join(formatted_names))
|
||||
|
||||
+683
-173
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,4 @@
|
||||
PREVIEW_CATEGORY = "🍃 FlowControl/Previews"
|
||||
PREVIEW_CATEGORY = "🪐 SaturnNodes/Previews"
|
||||
|
||||
class PreviewLatentLiveNode:
|
||||
@classmethod
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
from .prompt_iterator import parse_prompt_blocks
|
||||
|
||||
class PromptCounter:
|
||||
"""
|
||||
SaturnNodes Prompt Counter:
|
||||
Multiline text prompt input node with real-time prompt counting according to
|
||||
selectable separator ("Newline", ">1 Empty Line", ">2 Empty Lines", "Custom Regex").
|
||||
Outputs STRING (unmodified text) and count (integer count of prompts).
|
||||
"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"default": "", "multiline": True, "placeholder": "Enter multiline prompts here..."}),
|
||||
},
|
||||
"optional": {
|
||||
"separator": ([
|
||||
">1 Empty Line",
|
||||
"Newline",
|
||||
">2 Empty Lines",
|
||||
"Custom Regex"
|
||||
], {"default": ">1 Empty Line", "advanced": True}),
|
||||
"custom_regex": ("STRING", {"default": "", "advanced": True, "placeholder": "Custom regex delimiter (e.g. \\n---\\n)"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "INT")
|
||||
RETURN_NAMES = ("STRING", "count")
|
||||
FUNCTION = "count_prompts"
|
||||
CATEGORY = "🪐 SaturnNodes/Utils"
|
||||
DESCRIPTION = "Multiline prompt text input with live prompt count display according to selectable separator or custom regex."
|
||||
|
||||
def count_prompts(self, text="", separator=">1 Empty Line", custom_regex="", **kwargs):
|
||||
text_str = str(text) if text is not None else ""
|
||||
blocks = parse_prompt_blocks(text_str, separator, custom_regex)
|
||||
return (text_str, len(blocks))
|
||||
@@ -0,0 +1,519 @@
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
import copy
|
||||
import random
|
||||
import time
|
||||
import hashlib
|
||||
import platform
|
||||
import subprocess
|
||||
import threading
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
import folder_paths
|
||||
from .utils import get_saturnnodes_user_dir, is_authenticated_local_request, get_env_setting
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
STATE_FILE = os.path.join(USER_DIR, "prompt_iterator_state.json")
|
||||
|
||||
_STATE_LOCK = threading.RLock()
|
||||
_STATE_CACHE = None
|
||||
_STATE_CACHE_MTIME = None
|
||||
|
||||
def is_clear_prompt_iterator_on_launch():
|
||||
val = get_env_setting("CLEAR_PROMPT_ITERATOR_ON_LAUNCH", "false").lower()
|
||||
return val in ["true", "1", "yes"]
|
||||
|
||||
def ensure_user_dir():
|
||||
os.makedirs(USER_DIR, exist_ok=True)
|
||||
|
||||
def invalidate_state_cache():
|
||||
global _STATE_CACHE, _STATE_CACHE_MTIME
|
||||
with _STATE_LOCK:
|
||||
_STATE_CACHE = None
|
||||
_STATE_CACHE_MTIME = None
|
||||
|
||||
def load_state():
|
||||
global _STATE_CACHE, _STATE_CACHE_MTIME
|
||||
ensure_user_dir()
|
||||
with _STATE_LOCK:
|
||||
mtime = None
|
||||
try:
|
||||
if os.path.exists(STATE_FILE):
|
||||
mtime = os.stat(STATE_FILE).st_mtime_ns
|
||||
except OSError:
|
||||
mtime = None
|
||||
|
||||
if _STATE_CACHE is not None and mtime == _STATE_CACHE_MTIME:
|
||||
return copy.deepcopy(_STATE_CACHE)
|
||||
|
||||
data = {}
|
||||
if mtime is not None:
|
||||
try:
|
||||
with open(STATE_FILE, "r", encoding="utf-8") as f:
|
||||
loaded = json.load(f)
|
||||
if isinstance(loaded, dict):
|
||||
data = loaded
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Warning reading prompt iterator state: {e}")
|
||||
|
||||
_STATE_CACHE = copy.deepcopy(data)
|
||||
_STATE_CACHE_MTIME = mtime
|
||||
return copy.deepcopy(data)
|
||||
|
||||
def _entry_timestamp(entry):
|
||||
if isinstance(entry, dict):
|
||||
try:
|
||||
return int(entry.get("last_updated", 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
return 0
|
||||
|
||||
def save_state(state):
|
||||
global _STATE_CACHE, _STATE_CACHE_MTIME
|
||||
ensure_user_dir()
|
||||
with _STATE_LOCK:
|
||||
try:
|
||||
# Global prune: keep the 100 most recently touched node/text entries.
|
||||
if len(state) > 100:
|
||||
keep = set(sorted(state.keys(), key=lambda k: _entry_timestamp(state[k]), reverse=True)[:100])
|
||||
for k in [k for k in state.keys() if k not in keep]:
|
||||
del state[k]
|
||||
tmp_file = STATE_FILE + ".tmp"
|
||||
with open(tmp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(state, f, indent=4, ensure_ascii=False)
|
||||
os.replace(tmp_file, STATE_FILE)
|
||||
_STATE_CACHE = copy.deepcopy(state)
|
||||
try:
|
||||
_STATE_CACHE_MTIME = os.stat(STATE_FILE).st_mtime_ns
|
||||
except OSError:
|
||||
_STATE_CACHE_MTIME = None
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error saving prompt iterator state: {e}")
|
||||
|
||||
def clear_state():
|
||||
ensure_user_dir()
|
||||
with _STATE_LOCK:
|
||||
try:
|
||||
tmp_file = STATE_FILE + ".tmp"
|
||||
with open(tmp_file, "w", encoding="utf-8") as f:
|
||||
json.dump({}, f, indent=4)
|
||||
os.replace(tmp_file, STATE_FILE)
|
||||
invalidate_state_cache()
|
||||
print("[SaturnNodes] Prompt iterator state cleared.")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error clearing prompt iterator state: {e}")
|
||||
return False
|
||||
|
||||
# Clear state on startup if privacy setting enabled
|
||||
if is_clear_prompt_iterator_on_launch():
|
||||
clear_state()
|
||||
|
||||
# Register API endpoints
|
||||
try:
|
||||
routes = PromptServer.instance.routes
|
||||
|
||||
@routes.post("/saturnnodes/prompt_iterator/reset_node")
|
||||
async def reset_prompt_iterator_node(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
try:
|
||||
data = await request.json()
|
||||
node_id = str(data.get("node_id", ""))
|
||||
state = load_state()
|
||||
matched = reset_node_state(state, node_id)
|
||||
save_state(state)
|
||||
if node_id and node_id != "all":
|
||||
progress_payload = {
|
||||
"node_id": node_id,
|
||||
"current_run": 0,
|
||||
"total_runs": 0,
|
||||
"status_text": "Reset (0 / --)"
|
||||
}
|
||||
try:
|
||||
PromptServer.instance.send_sync("saturnnodes_prompt_iterator_progress", progress_payload)
|
||||
except Exception:
|
||||
pass
|
||||
return web.json_response({"status": "ok", "reset": matched})
|
||||
except Exception as e:
|
||||
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
||||
|
||||
@routes.post("/saturnnodes/prompt_iterator/clear")
|
||||
async def clear_prompt_iterator_endpoint(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
success = clear_state()
|
||||
return web.json_response({"status": "ok" if success else "error"})
|
||||
|
||||
@routes.post("/saturnnodes/prompt_iterator/open_file")
|
||||
async def open_prompt_iterator_file(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
ensure_user_dir()
|
||||
if not os.path.exists(STATE_FILE):
|
||||
save_state({})
|
||||
|
||||
# Security: verify canonical state file path is strictly within USER_DIR
|
||||
canonical_state = os.path.realpath(STATE_FILE)
|
||||
canonical_user = os.path.realpath(USER_DIR)
|
||||
try:
|
||||
if os.path.commonpath([canonical_user, canonical_state]) != canonical_user:
|
||||
return web.json_response({"status": "error", "message": "Security: State file escapes user directory"}, status=403)
|
||||
except Exception:
|
||||
return web.json_response({"status": "error", "message": "Security path validation error"}, status=403)
|
||||
|
||||
try:
|
||||
sys_name = platform.system()
|
||||
if sys_name == "Windows":
|
||||
os.startfile(canonical_state)
|
||||
elif sys_name == "Darwin":
|
||||
subprocess.Popen(["open", canonical_state])
|
||||
else:
|
||||
subprocess.Popen(["xdg-open", canonical_state])
|
||||
return web.json_response({"status": "ok", "path": canonical_state})
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error opening state file: {e}")
|
||||
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Notice: Skipped registering prompt iterator endpoints: {e}")
|
||||
|
||||
def parse_prompt_blocks(text_str, separator, custom_regex=""):
|
||||
if not text_str or not text_str.strip():
|
||||
return []
|
||||
|
||||
clean_text = text_str.replace('\r\n', '\n').replace('\r', '\n')
|
||||
sep_lower = str(separator).lower() if separator else ">1 empty line"
|
||||
|
||||
if "regex" in sep_lower:
|
||||
pattern = custom_regex.strip() if (custom_regex and custom_regex.strip()) else r'\n\s*\n+'
|
||||
try:
|
||||
regex = re.compile(pattern)
|
||||
raw_blocks = regex.split(clean_text)
|
||||
results = []
|
||||
for b in raw_blocks:
|
||||
b_str = b.strip()
|
||||
if b_str and not regex.fullmatch(b):
|
||||
results.append(b_str)
|
||||
return results
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes Prompt] Invalid custom regex pattern '{pattern}': {e}. Falling back to empty line separator.")
|
||||
raw_blocks = re.split(r'\n\s*\n+', clean_text)
|
||||
return [b.strip() for b in raw_blocks if b.strip()]
|
||||
elif "newline" in sep_lower:
|
||||
return [line.strip() for line in clean_text.split('\n') if line.strip()]
|
||||
elif ">2" in sep_lower or "2 empty" in sep_lower:
|
||||
raw_blocks = re.split(r'(?:\n\s*){3,}', clean_text)
|
||||
return [b.strip() for b in raw_blocks if b.strip()]
|
||||
else: # ">1 Empty Line", ">1 Emptyline"
|
||||
raw_blocks = re.split(r'\n\s*\n+', clean_text)
|
||||
return [b.strip() for b in raw_blocks if b.strip()]
|
||||
|
||||
def reset_node_state(state, node_id):
|
||||
"""Reset every pointer belonging to a node (or all nodes when node_id == 'all')."""
|
||||
node_id = str(node_id)
|
||||
matched = False
|
||||
for k in list(state.keys()):
|
||||
if node_id == "all" or k.startswith(f"node_{node_id}_") or k == f"node_{node_id}":
|
||||
if isinstance(state[k], dict):
|
||||
state[k]["index"] = 0
|
||||
state[k]["last_run"] = 0
|
||||
# A hard reset must not leave stale crash-recovery assignments or
|
||||
# shuffle permutations behind, otherwise the next retry can jump
|
||||
# to an old prompt or reuse an outdated cycle order.
|
||||
state[k]["assigned_prompts"] = {}
|
||||
state[k]["shuffled_order"] = []
|
||||
else:
|
||||
state[k] = {"index": 0, "total": 0}
|
||||
matched = True
|
||||
return matched
|
||||
|
||||
def get_current_prompt_id():
|
||||
try:
|
||||
server = getattr(PromptServer, "instance", None)
|
||||
if server and hasattr(server, "prompt_queue"):
|
||||
queue = server.prompt_queue
|
||||
if hasattr(queue, "currently_running"):
|
||||
mutex = getattr(queue, "mutex", None)
|
||||
if mutex:
|
||||
with mutex:
|
||||
items = list(queue.currently_running.values())
|
||||
else:
|
||||
items = list(queue.currently_running.values())
|
||||
for item in items:
|
||||
if isinstance(item, (tuple, list)) and len(item) > 1:
|
||||
return str(item[1])
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _remember_assignment(assigned_prompts, prompt_id, index, limit=50):
|
||||
"""Insert/refresh a crash-recovery mapping, keeping the most recent entries."""
|
||||
if not prompt_id:
|
||||
return
|
||||
assigned_prompts.pop(prompt_id, None)
|
||||
assigned_prompts[prompt_id] = index
|
||||
while len(assigned_prompts) > limit:
|
||||
assigned_prompts.pop(next(iter(assigned_prompts)))
|
||||
|
||||
def _sanitize_shuffle_order(raw_order, total_count):
|
||||
if not isinstance(raw_order, (list, tuple)):
|
||||
return None
|
||||
try:
|
||||
order = [int(v) for v in raw_order]
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if len(order) != total_count or len(set(order)) != total_count:
|
||||
return None
|
||||
if any(v < 0 or v >= total_count for v in order):
|
||||
return None
|
||||
return order
|
||||
|
||||
class PromptQueueIterator:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pop_mode": ([
|
||||
"Sequential (Loop on End)",
|
||||
"Sequential (Stop on End)",
|
||||
"Random (Keep)",
|
||||
"Random (Cycle)"
|
||||
], {"default": "Sequential (Loop on End)"}),
|
||||
"separator": ([
|
||||
">1 Empty Line",
|
||||
"Newline",
|
||||
">2 Empty Lines",
|
||||
"Custom Regex"
|
||||
], {"default": ">1 Empty Line"}),
|
||||
"text": ("STRING", {"default": "", "multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"custom_regex": ("STRING", {"default": "", "placeholder": "Custom regex delimiter (e.g. \\n---\\n)"}),
|
||||
# Appended after the existing widgets so old workflow widget
|
||||
# values keep their original positions.
|
||||
"track_by": ([
|
||||
"Text & Node",
|
||||
"Node Only"
|
||||
], {"default": "Text & Node"}),
|
||||
"prompt": ("STRING", {"forceInput": True}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "INT", "BOOLEAN")
|
||||
RETURN_NAMES = ("prompt", "remaining_text", "remaining_count", "exhausted")
|
||||
FUNCTION = "process_queue"
|
||||
CATEGORY = "🪐 SaturnNodes/Utils"
|
||||
DESCRIPTION = "Deterministically iterates over multiline prompt text blocks per queue run with live progress display and index controls."
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
# Force re-execution so each queue batch item advances the pointer
|
||||
return time.time()
|
||||
|
||||
def process_queue(
|
||||
self,
|
||||
pop_mode="Sequential (Loop on End)",
|
||||
track_by="Text & Node",
|
||||
separator=">1 Empty Line",
|
||||
text="",
|
||||
custom_regex="",
|
||||
prompt=None,
|
||||
prompt_text=None,
|
||||
unique_id="default",
|
||||
**kwargs
|
||||
):
|
||||
# A connected `prompt` input is authoritative. Falling back to the
|
||||
# local `text` widget used to replay stale characters from a previous
|
||||
# session even though the linked upstream node produced new prompts.
|
||||
if prompt is not None:
|
||||
raw_input = prompt
|
||||
elif text:
|
||||
raw_input = text
|
||||
else:
|
||||
raw_input = prompt_text
|
||||
text_str = str(raw_input) if raw_input is not None else ""
|
||||
if not text_str.strip():
|
||||
return ("", "", 0, False)
|
||||
|
||||
original_blocks = parse_prompt_blocks(text_str, separator, custom_regex)
|
||||
total_count = len(original_blocks)
|
||||
if total_count == 0:
|
||||
return ("", "", 0, False)
|
||||
|
||||
# Calculate stable hashes of the complete input text and of the splitting
|
||||
# configuration. Two nodes sharing text but using different separators must
|
||||
# never share a run pointer.
|
||||
text_hash = hashlib.sha256(text_str.strip().encode('utf-8')).hexdigest()[:16]
|
||||
config_signature = json.dumps({
|
||||
"mode": str(pop_mode),
|
||||
"separator": str(separator),
|
||||
"regex": str(custom_regex or ""),
|
||||
}, sort_keys=True)
|
||||
config_hash = hashlib.sha256(config_signature.encode('utf-8')).hexdigest()[:12]
|
||||
|
||||
track_lower = str(track_by or "").lower()
|
||||
if "node only" in track_lower:
|
||||
# Optional legacy-style tracking: one pointer per graph node, shared
|
||||
# across text edits and splitter changes.
|
||||
state_key = f"node_{unique_id}"
|
||||
else:
|
||||
state_key = f"node_{unique_id}_{config_hash}_{text_hash}"
|
||||
|
||||
current_prompt_id = get_current_prompt_id()
|
||||
|
||||
# Load state fresh from disk (no stale memory cache)
|
||||
state = load_state()
|
||||
node_state = state.get(state_key)
|
||||
|
||||
current_index = 0
|
||||
assigned_prompts = {}
|
||||
shuffled_order = []
|
||||
if isinstance(node_state, dict):
|
||||
try:
|
||||
current_index = max(0, int(node_state.get("index", 0)))
|
||||
except (TypeError, ValueError):
|
||||
current_index = 0
|
||||
raw_assigned = node_state.get("assigned_prompts", {})
|
||||
if isinstance(raw_assigned, dict):
|
||||
assigned_prompts = dict(raw_assigned)
|
||||
sanitized = _sanitize_shuffle_order(node_state.get("shuffled_order", []), total_count)
|
||||
if sanitized is not None:
|
||||
shuffled_order = sanitized
|
||||
elif isinstance(node_state, int):
|
||||
current_index = max(0, node_state)
|
||||
|
||||
is_stop_mode = pop_mode == "Sequential (Stop on End)"
|
||||
is_random_keep = pop_mode in ["Random (Keep)", "Random (Delete)"]
|
||||
is_random_cycle = pop_mode in ["Random (Cycle)"]
|
||||
|
||||
if is_random_cycle and not shuffled_order:
|
||||
shuffled_order = list(range(total_count))
|
||||
random.shuffle(shuffled_order)
|
||||
|
||||
exhausted = False
|
||||
remaining_count = 0
|
||||
remaining_text = ""
|
||||
|
||||
# Check if this prompt_id was already assigned an index (e.g. recovered from persistent queue / retry)
|
||||
retry_assignment = current_prompt_id and current_prompt_id in assigned_prompts
|
||||
if retry_assignment:
|
||||
try:
|
||||
idx = int(assigned_prompts[current_prompt_id]) % total_count
|
||||
except (TypeError, ValueError):
|
||||
idx = current_index % total_count
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = current_index # Keep current index without double-advancing
|
||||
# Refresh recency so long queues do not evict actively retried prompts.
|
||||
_remember_assignment(assigned_prompts, current_prompt_id, idx)
|
||||
else:
|
||||
# Select prompt based on mode and calculate next index
|
||||
if pop_mode in ["Sequential (Loop on End)", "Cycle / Loop", "Pop Top & Delete"]:
|
||||
idx = current_index % total_count
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = (current_index + 1) % total_count
|
||||
elif is_stop_mode:
|
||||
if current_index >= total_count:
|
||||
idx = total_count - 1
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = total_count
|
||||
else:
|
||||
idx = current_index
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = current_index + 1
|
||||
elif is_random_keep:
|
||||
idx = random.randint(0, total_count - 1)
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = current_index + 1
|
||||
elif is_random_cycle:
|
||||
cycle_index = current_index % total_count
|
||||
if len(shuffled_order) != total_count or (cycle_index == 0 and current_index > 0):
|
||||
shuffled_order = list(range(total_count))
|
||||
random.shuffle(shuffled_order)
|
||||
idx = shuffled_order[cycle_index]
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = (current_index + 1) % total_count
|
||||
else: # Default fallback
|
||||
idx = current_index % total_count
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = (current_index + 1) % total_count
|
||||
|
||||
_remember_assignment(assigned_prompts, current_prompt_id, idx)
|
||||
|
||||
# Calculate progress and the real remaining set for every mode.
|
||||
if is_random_cycle:
|
||||
try:
|
||||
cycle_position = shuffled_order.index(idx)
|
||||
except ValueError:
|
||||
cycle_position = idx
|
||||
display_run = cycle_position + 1
|
||||
remaining_blocks = [original_blocks[i] for i in shuffled_order[cycle_position + 1:]]
|
||||
remaining_count = len(remaining_blocks)
|
||||
elif is_random_keep:
|
||||
# Random (Keep) never consumes prompts and has no deterministic order.
|
||||
display_run = idx + 1 if retry_assignment else (current_index % total_count) + 1
|
||||
remaining_blocks = []
|
||||
remaining_count = 0
|
||||
else:
|
||||
display_run = idx + 1
|
||||
remaining_blocks = original_blocks[idx + 1:]
|
||||
remaining_count = len(remaining_blocks)
|
||||
|
||||
if is_stop_mode and remaining_count == 0:
|
||||
exhausted = True
|
||||
|
||||
sep_lower = str(separator).lower() if separator else ""
|
||||
if "newline" in sep_lower:
|
||||
join_delim = "\n"
|
||||
elif ">2" in sep_lower or "2 empty" in sep_lower:
|
||||
join_delim = "\n\n\n"
|
||||
elif "regex" in sep_lower:
|
||||
if custom_regex and not any(c in custom_regex for c in r"*+?{}[\]^$|()\\"):
|
||||
join_delim = custom_regex
|
||||
else:
|
||||
join_delim = "\n\n"
|
||||
else:
|
||||
join_delim = "\n\n"
|
||||
remaining_text = join_delim.join(remaining_blocks)
|
||||
|
||||
# Update state dictionary
|
||||
new_entry = {
|
||||
"index": next_index,
|
||||
"total": total_count,
|
||||
"last_run": display_run,
|
||||
"last_updated": int(time.time()),
|
||||
"preview": selected_prompt[:60].replace('\n', ' '),
|
||||
"assigned_prompts": assigned_prompts
|
||||
}
|
||||
if is_random_cycle:
|
||||
new_entry["shuffled_order"] = shuffled_order
|
||||
state[state_key] = new_entry
|
||||
|
||||
# Keep at most 3 most recent text hashes/configs per node_id
|
||||
node_prefix = f"node_{unique_id}_"
|
||||
node_keys = [k for k in state.keys() if k.startswith(node_prefix)]
|
||||
if len(node_keys) > 3:
|
||||
node_keys.sort(key=lambda k: _entry_timestamp(state[k]))
|
||||
for old_k in node_keys[:-3]:
|
||||
del state[old_k]
|
||||
|
||||
save_state(state)
|
||||
|
||||
# Send live progress update to UI
|
||||
progress_data = {
|
||||
"node_id": str(unique_id),
|
||||
"current_run": display_run,
|
||||
"total_runs": total_count,
|
||||
"status_text": f"Run {display_run} / {total_count}"
|
||||
}
|
||||
try:
|
||||
PromptServer.instance.send_sync("saturnnodes_prompt_iterator_progress", progress_data)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return (selected_prompt, remaining_text, remaining_count, exhausted)
|
||||
+989
-187
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,476 @@
|
||||
import os
|
||||
import urllib.request
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
folder_paths = None
|
||||
|
||||
# Global ONNX runtime session cache & predicted depth cache
|
||||
_DEPTH_ORT_SESSION = None
|
||||
_PREDICTED_DEPTH_CACHE = {}
|
||||
|
||||
def get_depth_model_path():
|
||||
"""Returns local path to Depth Anything V2 Small ONNX model."""
|
||||
if folder_paths and hasattr(folder_paths, "models_dir") and folder_paths.models_dir:
|
||||
base_dir = os.path.join(folder_paths.models_dir, "depth")
|
||||
else:
|
||||
base_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models")
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
return os.path.join(base_dir, "depth-anything-v2-small-518.onnx")
|
||||
|
||||
def ensure_depth_model():
|
||||
"""Downloads Depth Anything V2 Small ONNX model if not already present."""
|
||||
target_path = get_depth_model_path()
|
||||
if os.path.exists(target_path) and os.path.getsize(target_path) > 10 * 1024 * 1024:
|
||||
return target_path
|
||||
|
||||
temp_path = target_path + ".tmp"
|
||||
urls = [
|
||||
"https://huggingface.co/onnx-community/depth-anything-v2-small/resolve/main/onnx/model.onnx",
|
||||
"https://nqrlabs.com/Stereonix/models/depth-anything-v2-small-518.onnx"
|
||||
]
|
||||
|
||||
print("[SaturnNodes 3D] Downloading Depth Anything V2 Small model (~94 MB)...")
|
||||
for url in urls:
|
||||
try:
|
||||
req = urllib.request.Request(url, headers={"User-Agent": "ComfyUI-SaturnNodes"})
|
||||
with urllib.request.urlopen(req, timeout=60) as resp, open(temp_path, "wb") as f:
|
||||
total_size = int(resp.headers.get("Content-Length", 0))
|
||||
downloaded = 0
|
||||
last_pct = 0
|
||||
while True:
|
||||
chunk = resp.read(1024 * 1024 * 2)
|
||||
if not chunk:
|
||||
break
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
if total_size > 0:
|
||||
pct = int((downloaded / total_size) * 100)
|
||||
if pct >= last_pct + 25 or pct == 100:
|
||||
print(f"[SaturnNodes 3D] Model download: {pct}% ({downloaded // (1024*1024)}MB / {total_size // (1024*1024)}MB)")
|
||||
last_pct = pct
|
||||
|
||||
if os.path.exists(temp_path) and os.path.getsize(temp_path) > 10 * 1024 * 1024:
|
||||
os.replace(temp_path, target_path)
|
||||
print(f"[SaturnNodes 3D] Successfully downloaded Depth Anything V2 to {target_path}")
|
||||
return target_path
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes 3D] Download error from {url}: {e}")
|
||||
if os.path.exists(temp_path):
|
||||
try:
|
||||
os.remove(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
raise RuntimeError("Failed to download Depth Anything V2 Small ONNX model.")
|
||||
|
||||
def get_ort_session():
|
||||
"""Retrieves or loads the ONNX runtime session for depth estimation."""
|
||||
global _DEPTH_ORT_SESSION
|
||||
if _DEPTH_ORT_SESSION is not None:
|
||||
return _DEPTH_ORT_SESSION
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
except ImportError:
|
||||
raise RuntimeError("onnxruntime is required for built-in depth estimation. Install via: pip install onnxruntime")
|
||||
|
||||
model_path = ensure_depth_model()
|
||||
available_providers = ort.get_available_providers()
|
||||
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in available_providers else ["CPUExecutionProvider"]
|
||||
|
||||
_DEPTH_ORT_SESSION = ort.InferenceSession(model_path, providers=providers)
|
||||
return _DEPTH_ORT_SESSION
|
||||
|
||||
def predict_depth_map(image_np):
|
||||
"""
|
||||
Runs Depth Anything V2 Small on a single [H, W, 3] float32 image (0.0 to 1.0).
|
||||
Returns [H, W] normalized float32 depth map (0.0=far, 1.0=near).
|
||||
If onnxruntime or the model is unavailable, gracefully falls back to a pseudo-depth map.
|
||||
"""
|
||||
H, W, _ = image_np.shape
|
||||
try:
|
||||
session = get_ort_session()
|
||||
if session is None:
|
||||
raise RuntimeError("ONNX runtime session could not be initialized.")
|
||||
|
||||
target_size = 518
|
||||
scale = min(target_size / W, target_size / H)
|
||||
scaled_w = max(1, int(round(W * scale)))
|
||||
scaled_h = max(1, int(round(H * scale)))
|
||||
offset_x = (target_size - scaled_w) // 2
|
||||
offset_y = (target_size - scaled_h) // 2
|
||||
|
||||
# Letterbox to 518x518 with black padding using PIL
|
||||
from PIL import Image
|
||||
uint_img = (image_np * 255.0).clip(0, 255).astype(np.uint8)
|
||||
scaled_pil = Image.fromarray(uint_img).resize((scaled_w, scaled_h), Image.Resampling.BILINEAR)
|
||||
scaled_img = np.array(scaled_pil, dtype=np.uint8)
|
||||
|
||||
canvas = np.zeros((target_size, target_size, 3), dtype=np.uint8)
|
||||
canvas[offset_y:offset_y + scaled_h, offset_x:offset_x + scaled_w] = scaled_img
|
||||
|
||||
# Normalize with ImageNet constants
|
||||
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
|
||||
std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
|
||||
norm = ((canvas.astype(np.float32) / 255.0) - mean) / std
|
||||
input_tensor = np.transpose(norm, (2, 0, 1))[np.newaxis, ...].astype(np.float32)
|
||||
|
||||
# Run inference
|
||||
input_name = session.get_inputs()[0].name
|
||||
raw_depth = session.run(None, {input_name: input_tensor})[0][0] # [518, 518]
|
||||
|
||||
# Crop out the letterbox padding and resize back to original (W, H)
|
||||
valid_depth = raw_depth[offset_y:offset_y + scaled_h, offset_x:offset_x + scaled_w]
|
||||
depth_pil = Image.fromarray(valid_depth.astype(np.float32), mode="F").resize((W, H), Image.Resampling.BILINEAR)
|
||||
depth_full = np.array(depth_pil, dtype=np.float32)
|
||||
|
||||
# Normalize with p2 and p98 histogram percentiles
|
||||
p2 = np.percentile(depth_full, 2)
|
||||
p98 = np.percentile(depth_full, 98)
|
||||
denom = max(float(p98 - p2), 1e-6)
|
||||
depth_norm = np.clip((depth_full - p2) / denom, 0.0, 1.0).astype(np.float32)
|
||||
return depth_norm
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes 3D] Notice: Automatic depth model unavailable ({e}). Using pseudo-depth gradient fallback.")
|
||||
y_gradient = np.linspace(0.2, 1.0, H, dtype=np.float32).reshape(H, 1)
|
||||
luminance = np.dot(image_np[..., :3], [0.299, 0.587, 0.114])
|
||||
fallback = 0.7 * y_gradient + 0.3 * luminance
|
||||
denom = max(float(fallback.max() - fallback.min()), 1e-6)
|
||||
return ((fallback - fallback.min()) / denom).astype(np.float32)
|
||||
|
||||
def draw_alignment_dot(target_tensor, cx, cy, radius=5):
|
||||
"""Draws a high-contrast white dot with black outline at (cx, cy) on [H, W, 3] tensor."""
|
||||
H, W, _ = target_tensor.shape
|
||||
y_min = max(0, cy - radius - 2)
|
||||
y_max = min(H, cy + radius + 3)
|
||||
x_min = max(0, cx - radius - 2)
|
||||
x_max = min(W, cx + radius + 3)
|
||||
|
||||
if y_max <= y_min or x_max <= x_min:
|
||||
return
|
||||
|
||||
yy, xx = torch.meshgrid(
|
||||
torch.arange(y_min, y_max, device=target_tensor.device),
|
||||
torch.arange(x_min, x_max, device=target_tensor.device),
|
||||
indexing="ij"
|
||||
)
|
||||
dist_sq = (xx - cx) ** 2 + (yy - cy) ** 2
|
||||
inner = dist_sq <= (radius ** 2)
|
||||
outer = (dist_sq <= ((radius + 1.5) ** 2)) & (~inner)
|
||||
|
||||
target_tensor[yy[outer], xx[outer]] = 0.0
|
||||
target_tensor[yy[inner], xx[inner]] = 1.0
|
||||
|
||||
|
||||
class SaturnStereo3D:
|
||||
"""
|
||||
🪐 Stereoscopic 3D / Cross-Eye Generator
|
||||
Generates side-by-side 3D stereograms, parallel views, and red-cyan anaglyphs
|
||||
from 2D images using depth-based horizontal disparity mapping.
|
||||
Includes built-in Depth Anything V2 Small auto-inference, zoom & pan cropping,
|
||||
spacing/border padding, and alignment guide dots.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"mode": ([
|
||||
"Cross-Eye",
|
||||
"Parallel (Wall-Eyed)",
|
||||
"Red-Cyan Anaglyph",
|
||||
"Side-by-Side (Parallel)",
|
||||
"Top-Bottom (Over-Under)",
|
||||
"Left Eye Only",
|
||||
"Right Eye Only"
|
||||
], {"default": "Cross-Eye"}),
|
||||
"depth_intensity": ("FLOAT", {
|
||||
"default": 25.0, "min": 0.0, "max": 100.0, "step": 1.0,
|
||||
"display": "slider",
|
||||
"tooltip": "Controls how much the left and right perspectives shift apart."
|
||||
}),
|
||||
"depth_gamma": ("FLOAT", {
|
||||
"default": 0.50, "min": 0.1, "max": 3.0, "step": 0.05,
|
||||
"display": "slider",
|
||||
"tooltip": "Adjust depth curve for realistic 3D falloff (Stereonix default is 0.50)."
|
||||
}),
|
||||
"invert_depth": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Flip which parts appear closer or further away."
|
||||
}),
|
||||
"convergence_plane": ("FLOAT", {
|
||||
"default": 0.50, "min": 0.0, "max": 1.0, "step": 0.02,
|
||||
"display": "slider",
|
||||
"tooltip": "Depth level where parallax is zero (screen surface). Above pops out, below sinks in."
|
||||
}),
|
||||
"image_zoom": ("FLOAT", {
|
||||
"default": 100.0, "min": 10.0, "max": 300.0, "step": 1.0,
|
||||
"display": "slider",
|
||||
"tooltip": "Zoom into areas of interest before generating 3D disparity."
|
||||
}),
|
||||
"pan_horizontal": ("FLOAT", {
|
||||
"default": 0.0, "min": -100.0, "max": 100.0, "step": 1.0,
|
||||
"display": "slider",
|
||||
"tooltip": "Horizontal pan offset (-100 to 100)."
|
||||
}),
|
||||
"pan_vertical": ("FLOAT", {
|
||||
"default": 0.0, "min": -100.0, "max": 100.0, "step": 1.0,
|
||||
"display": "slider",
|
||||
"tooltip": "Vertical pan offset (-100 to 100)."
|
||||
}),
|
||||
"gap_spacing": ("INT", {
|
||||
"default": 20, "min": 0, "max": 1000, "step": 2,
|
||||
"display": "number",
|
||||
"tooltip": "Empty black space between the left and right images."
|
||||
}),
|
||||
"outer_border": ("INT", {
|
||||
"default": 0, "min": 0, "max": 1000, "step": 2,
|
||||
"display": "number",
|
||||
"tooltip": "Black border padding around the canvas (allows shrinking images for easier cross-eye viewing)."
|
||||
}),
|
||||
"alignment_dots": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Place focus guide dots above each eye view to easily lock into 3D cross-eye view."
|
||||
}),
|
||||
"fill_method": ([
|
||||
"Smooth (Bilinear Grid)",
|
||||
"Nearest Neighbor (Stereonix)"
|
||||
], {
|
||||
"default": "Smooth (Bilinear Grid)",
|
||||
"tooltip": "Hole-filling interpolation for occlusion gaps."
|
||||
}),
|
||||
"preview_zoom": ([
|
||||
"Fit",
|
||||
"25%",
|
||||
"33%",
|
||||
"50%",
|
||||
"75%",
|
||||
"100%"
|
||||
], {
|
||||
"default": "50%",
|
||||
"tooltip": "Scale preview for easier cross-eye viewing on canvas (does not affect full output)."
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"depth_map": ("IMAGE", {
|
||||
"tooltip": "Optional depth map. If disconnected, automatically runs Depth Anything V2."
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE")
|
||||
RETURN_NAMES = ("stereogram", "left_eye", "right_eye", "depth_map")
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "🪐 SaturnNodes/3D & Previews"
|
||||
DESCRIPTION = (
|
||||
"Creates side-by-side cross-eye stereograms, parallel 3D views, and red-cyan anaglyphs "
|
||||
"from 2D images. Inspired by Stereonix, it includes automatic Depth Anything V2 Small fallback, "
|
||||
"custom depth gamma/inversion, zoom & pan cropping, adjustable gap/border spacing, and focus alignment dots."
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(
|
||||
self,
|
||||
image,
|
||||
mode="Cross-Eye",
|
||||
depth_intensity=25.0,
|
||||
depth_gamma=0.50,
|
||||
invert_depth=False,
|
||||
convergence_plane=0.50,
|
||||
image_zoom=100.0,
|
||||
pan_horizontal=0.0,
|
||||
pan_vertical=0.0,
|
||||
gap_spacing=20,
|
||||
outer_border=0,
|
||||
alignment_dots=True,
|
||||
fill_method="Smooth (Bilinear Grid)",
|
||||
preview_zoom="50%",
|
||||
depth_map=None,
|
||||
**kwargs
|
||||
):
|
||||
device = image.device
|
||||
B, H, W, C = image.shape
|
||||
|
||||
# 1. Obtain or generate depth map
|
||||
depth_tensors = []
|
||||
if depth_map is not None:
|
||||
# Handle incoming depth map tensor
|
||||
d_map = depth_map.to(device=device, dtype=torch.float32)
|
||||
if d_map.shape[1:3] != (H, W):
|
||||
d_perm = d_map.permute(0, 3, 1, 2) if d_map.dim() == 4 else d_map.unsqueeze(1)
|
||||
d_resized = F.interpolate(d_perm, size=(H, W), mode="bilinear", align_corners=False)
|
||||
d_map = d_resized.permute(0, 2, 3, 1) if d_map.dim() == 4 else d_resized.squeeze(1)
|
||||
|
||||
if d_map.dim() == 4:
|
||||
d_tensor = d_map[..., 0] # take single channel
|
||||
else:
|
||||
d_tensor = d_map
|
||||
|
||||
# Normalize 0 to 1
|
||||
for b in range(B):
|
||||
b_depth = d_tensor[b if b < d_tensor.shape[0] else 0]
|
||||
min_v = b_depth.min()
|
||||
max_v = b_depth.max()
|
||||
denom = max(float(max_v - min_v), 1e-6)
|
||||
depth_norm = torch.clamp((b_depth - min_v) / denom, 0.0, 1.0)
|
||||
depth_tensors.append(depth_norm)
|
||||
else:
|
||||
# Automatic Depth Estimation via Depth Anything V2 Small ONNX (cached for instant slider responsiveness)
|
||||
cpu_img = image.detach().cpu().numpy()
|
||||
for b in range(B):
|
||||
# Fast signature based on shape, sum, and corner samples
|
||||
img_slice = cpu_img[b]
|
||||
cache_key = (
|
||||
img_slice.shape,
|
||||
round(float(img_slice.sum()), 2),
|
||||
round(float(img_slice[0, 0, 0]), 4),
|
||||
round(float(img_slice[-1, -1, -1]), 4)
|
||||
)
|
||||
if cache_key in _PREDICTED_DEPTH_CACHE:
|
||||
d_np = _PREDICTED_DEPTH_CACHE[cache_key]
|
||||
else:
|
||||
d_np = predict_depth_map(img_slice)
|
||||
while len(_PREDICTED_DEPTH_CACHE) >= 16:
|
||||
_PREDICTED_DEPTH_CACHE.pop(next(iter(_PREDICTED_DEPTH_CACHE)))
|
||||
_PREDICTED_DEPTH_CACHE[cache_key] = d_np
|
||||
depth_tensors.append(torch.from_numpy(d_np).to(device=device, dtype=torch.float32))
|
||||
|
||||
depth_batch = torch.stack(depth_tensors, dim=0) # [B, H, W]
|
||||
|
||||
# 2. Apply Invert & Depth Gamma
|
||||
if invert_depth:
|
||||
depth_batch = 1.0 - depth_batch
|
||||
|
||||
if abs(depth_gamma - 1.0) > 1e-4:
|
||||
depth_batch = torch.pow(torch.clamp(depth_batch, 1e-6, 1.0), depth_gamma)
|
||||
|
||||
# 3. Apply Zoom & Pan cropping if requested
|
||||
zoom_factor = max(0.1, float(image_zoom) / 100.0)
|
||||
pan_x_norm = float(pan_horizontal) / 100.0
|
||||
pan_y_norm = float(pan_vertical) / 100.0
|
||||
|
||||
grid_y = torch.linspace(-1, 1, H, device=device).view(1, H, 1, 1).expand(B, H, W, 1)
|
||||
grid_x = torch.linspace(-1, 1, W, device=device).view(1, 1, W, 1).expand(B, H, W, 1)
|
||||
|
||||
if abs(zoom_factor - 1.0) > 1e-4 or abs(pan_x_norm) > 1e-4 or abs(pan_y_norm) > 1e-4:
|
||||
crop_x = (grid_x - pan_x_norm) / zoom_factor
|
||||
crop_y = (grid_y - pan_y_norm) / zoom_factor
|
||||
crop_grid = torch.cat([crop_x, crop_y], dim=-1)
|
||||
|
||||
img_perm = image.permute(0, 3, 1, 2)
|
||||
d_perm = depth_batch.unsqueeze(1)
|
||||
image_cropped = F.grid_sample(img_perm, crop_grid, mode="bilinear", padding_mode="zeros", align_corners=True).permute(0, 2, 3, 1)
|
||||
depth_cropped = F.grid_sample(d_perm, crop_grid, mode="bilinear", padding_mode="border", align_corners=True).squeeze(1)
|
||||
else:
|
||||
image_cropped = image
|
||||
depth_cropped = depth_batch
|
||||
|
||||
# 4. Compute Parallax Shift (Horizontal Disparity)
|
||||
# Shift magnitude scales with image width: at intensity 100, maximum shift is 3.5% of width
|
||||
max_shift = (float(depth_intensity) / 100.0) * (W * 0.035)
|
||||
# Shift relative to convergence plane
|
||||
shift_px = (depth_cropped - float(convergence_plane)) * max_shift
|
||||
shift_norm = (shift_px * (2.0 / W)).unsqueeze(-1) # [B, H, W, 1]
|
||||
|
||||
# 5. Generate Left and Right eye views
|
||||
grid_left = torch.cat([grid_x - shift_norm, grid_y], dim=-1)
|
||||
grid_right = torch.cat([grid_x + shift_norm, grid_y], dim=-1)
|
||||
|
||||
img_in = image_cropped.permute(0, 3, 1, 2)
|
||||
|
||||
if fill_method == "Nearest Neighbor (Stereonix)":
|
||||
# Backward warp with nearest-neighbor lookup matching classic Stereonix
|
||||
left_eye = F.grid_sample(img_in, grid_left, mode="nearest", padding_mode="border", align_corners=True).permute(0, 2, 3, 1)
|
||||
right_eye = F.grid_sample(img_in, grid_right, mode="nearest", padding_mode="border", align_corners=True).permute(0, 2, 3, 1)
|
||||
else:
|
||||
# Smooth bilinear grid sampling
|
||||
left_eye = F.grid_sample(img_in, grid_left, mode="bilinear", padding_mode="border", align_corners=True).permute(0, 2, 3, 1)
|
||||
right_eye = F.grid_sample(img_in, grid_right, mode="bilinear", padding_mode="border", align_corners=True).permute(0, 2, 3, 1)
|
||||
|
||||
# 6. Composite according to Selected Mode
|
||||
if mode == "Left Eye Only":
|
||||
stereogram = left_eye
|
||||
elif mode == "Right Eye Only":
|
||||
stereogram = right_eye
|
||||
elif mode == "Red-Cyan Anaglyph":
|
||||
# Red channel from Left Eye, Cyan (Green + Blue) channels from Right Eye
|
||||
stereogram = torch.zeros_like(left_eye)
|
||||
stereogram[..., 0] = left_eye[..., 0]
|
||||
stereogram[..., 1] = right_eye[..., 1]
|
||||
stereogram[..., 2] = right_eye[..., 2]
|
||||
elif mode == "Top-Bottom (Over-Under)":
|
||||
gap_h = max(0, int(gap_spacing))
|
||||
border = max(0, int(outer_border))
|
||||
|
||||
top_panel = left_eye
|
||||
bot_panel = right_eye
|
||||
|
||||
if gap_h > 0:
|
||||
gap_row = torch.zeros((B, gap_h, W, C), device=device, dtype=torch.float32)
|
||||
stacked = torch.cat([top_panel, gap_row, bot_panel], dim=1)
|
||||
else:
|
||||
stacked = torch.cat([top_panel, bot_panel], dim=1)
|
||||
|
||||
if border > 0:
|
||||
H_st, W_st = stacked.shape[1], stacked.shape[2]
|
||||
stereogram = torch.zeros((B, H_st + border * 2, W_st + border * 2, C), device=device, dtype=torch.float32)
|
||||
stereogram[:, border:border + H_st, border:border + W_st, :] = stacked
|
||||
else:
|
||||
stereogram = stacked
|
||||
else:
|
||||
# Side-by-Side: Cross-Eye vs Parallel
|
||||
# In Cross-Eye: Left panel is Right-Eye, Right panel is Left-Eye
|
||||
# In Parallel / Side-by-Side: Left panel is Left-Eye, Right panel is Right-Eye
|
||||
if mode == "Cross-Eye":
|
||||
panel_left = right_eye
|
||||
panel_right = left_eye
|
||||
else:
|
||||
panel_left = left_eye
|
||||
panel_right = right_eye
|
||||
|
||||
gap_w = max(0, int(gap_spacing))
|
||||
border = max(0, int(outer_border))
|
||||
|
||||
if gap_w > 0:
|
||||
gap_col = torch.zeros((B, H, gap_w, C), device=device, dtype=torch.float32)
|
||||
side_by_side = torch.cat([panel_left, gap_col, panel_right], dim=2)
|
||||
else:
|
||||
side_by_side = torch.cat([panel_left, panel_right], dim=2)
|
||||
|
||||
# Apply outer border framing
|
||||
if border > 0:
|
||||
H_sbs, W_sbs = side_by_side.shape[1], side_by_side.shape[2]
|
||||
stereogram = torch.zeros((B, H_sbs + border * 2, W_sbs + border * 2, C), device=device, dtype=torch.float32)
|
||||
stereogram[:, border:border + H_sbs, border:border + W_sbs, :] = side_by_side
|
||||
else:
|
||||
stereogram = side_by_side
|
||||
|
||||
# Add alignment guide dots
|
||||
if alignment_dots and (mode in ("Cross-Eye", "Parallel (Wall-Eyed)", "Side-by-Side (Parallel)")):
|
||||
# Center coordinates above left and right panels
|
||||
cx_left = border + (W // 2)
|
||||
cx_right = border + W + gap_w + (W // 2)
|
||||
cy_dot = border // 2 if border >= 18 else min(16, H // 12)
|
||||
|
||||
dot_radius = max(4, min(8, W // 100))
|
||||
for b in range(B):
|
||||
draw_alignment_dot(stereogram[b], cx_left, cy_dot, radius=dot_radius)
|
||||
draw_alignment_dot(stereogram[b], cx_right, cy_dot, radius=dot_radius)
|
||||
|
||||
# 7. Convert Depth to 3-channel image for output
|
||||
depth_output = depth_cropped.unsqueeze(-1).repeat(1, 1, 1, 3)
|
||||
|
||||
try:
|
||||
import comfy.model_management
|
||||
comfy.model_management.soft_empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return (stereogram, left_eye, right_eye, depth_output)
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
import re
|
||||
import csv
|
||||
import io
|
||||
|
||||
def parse_find_targets(find_str):
|
||||
if not find_str or not find_str.strip():
|
||||
return []
|
||||
try:
|
||||
reader = csv.reader(io.StringIO(find_str.strip()), skipinitialspace=True)
|
||||
raw_items = []
|
||||
for row in reader:
|
||||
for item in row:
|
||||
s = item.strip()
|
||||
if (s.startswith('"') and s.endswith('"')) or (s.startswith("'") and s.endswith("'")):
|
||||
s = s[1:-1].strip()
|
||||
if s:
|
||||
raw_items.append(s)
|
||||
return raw_items
|
||||
except Exception:
|
||||
raw_items = []
|
||||
parts = find_str.split(',')
|
||||
for part in parts:
|
||||
item = part.strip()
|
||||
if (item.startswith('"') and item.endswith('"')) or (item.startswith("'") and item.endswith("'")):
|
||||
item = item[1:-1].strip()
|
||||
if item:
|
||||
raw_items.append(item)
|
||||
return raw_items
|
||||
|
||||
class MultiTextReplacer:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"find": ("STRING", {"default": "", "multiline": True}),
|
||||
"replace": ("STRING", {"default": "", "multiline": False}),
|
||||
"case_sensitive": ("BOOLEAN", {"default": False, "advanced": True}),
|
||||
"search_mode": (["Comma Separated List", "Regex Pattern"], {"default": "Comma Separated List", "advanced": True}),
|
||||
},
|
||||
"optional": {
|
||||
"text": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "INT")
|
||||
RETURN_NAMES = ("text", "replaced_count")
|
||||
FUNCTION = "replace_text"
|
||||
CATEGORY = "🪐 SaturnNodes/Utils"
|
||||
DESCRIPTION = "Searches input text for multiple search targets specified in a comma-separated list or regex, and replaces all matches with a replacement string."
|
||||
|
||||
def replace_text(
|
||||
self,
|
||||
find="",
|
||||
replace="",
|
||||
case_sensitive=False,
|
||||
search_mode="Comma Separated List",
|
||||
text=None
|
||||
):
|
||||
input_text = str(text) if text is not None else ""
|
||||
if not input_text or not find:
|
||||
return (input_text, 0)
|
||||
|
||||
flags = 0 if case_sensitive else re.IGNORECASE
|
||||
replaced_count = 0
|
||||
result_text = input_text
|
||||
|
||||
if search_mode == "Comma Separated List":
|
||||
targets = parse_find_targets(find)
|
||||
if not targets:
|
||||
return (input_text, 0)
|
||||
|
||||
# Sort targets by length descending so longer phrases match first
|
||||
targets.sort(key=len, reverse=True)
|
||||
|
||||
# Build single unified regex pattern to guarantee 100% protection against recursive replacement loops
|
||||
escaped_targets = [re.escape(t) for t in targets]
|
||||
unified_pattern = "(?:" + "|".join(escaped_targets) + ")"
|
||||
|
||||
# Count total matches
|
||||
replaced_count = len(re.findall(unified_pattern, result_text, flags=flags))
|
||||
# Perform single-pass substitution
|
||||
result_text = re.sub(unified_pattern, lambda m: replace, result_text, flags=flags)
|
||||
else: # Regex Pattern
|
||||
lines = [line.strip() for line in find.splitlines() if line.strip()]
|
||||
for pattern_str in lines:
|
||||
try:
|
||||
count = len(re.findall(pattern_str, result_text, flags=flags))
|
||||
replaced_count += count
|
||||
result_text = re.sub(pattern_str, lambda m: replace, result_text, flags=flags)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Invalid regex pattern '{pattern_str}': {e}")
|
||||
|
||||
return (result_text, replaced_count)
|
||||
@@ -0,0 +1,87 @@
|
||||
import re
|
||||
|
||||
class SaturnTextSplit:
|
||||
"""
|
||||
SaturnNodes Text Split Node:
|
||||
Splits input text into two parts (text1 and text2) by either a literal character sequence
|
||||
or a regular expression pattern, supporting forward (from start) and backward (from end) splitting.
|
||||
"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"default": "", "multiline": True, "placeholder": "Text to split..."}),
|
||||
"split_by": ("STRING", {"default": "--", "multiline": False, "placeholder": "Character sequence or regex pattern"}),
|
||||
},
|
||||
"optional": {
|
||||
"use_regex": ("BOOLEAN", {"default": False, "tooltip": "Interpret 'split_by' as a Regular Expression instead of a literal string."}),
|
||||
"split_direction": (["forward (first occurrence from start)", "backward (last occurrence from end)"], {"default": "forward (first occurrence from start)", "advanced": True, "tooltip": "Choose whether to split on the first match from the beginning or the last match from the end."}),
|
||||
"strip_whitespace": ("BOOLEAN", {"default": False, "advanced": True, "tooltip": "Strip leading/trailing whitespace from the resulting text1 and text2."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING")
|
||||
RETURN_NAMES = ("text1", "text2")
|
||||
FUNCTION = "split_text"
|
||||
CATEGORY = "🪐 SaturnNodes/Utils"
|
||||
DESCRIPTION = "Splits input text into two parts (text1 and text2) by a character sequence or regex pattern, with forward/backward splitting."
|
||||
|
||||
def split_text(
|
||||
self,
|
||||
text="",
|
||||
split_by="--",
|
||||
use_regex=False,
|
||||
split_direction="forward (first occurrence from start)",
|
||||
strip_whitespace=False
|
||||
):
|
||||
input_text = str(text) if text is not None else ""
|
||||
delimiter = str(split_by) if split_by is not None else ""
|
||||
|
||||
if not input_text or not delimiter:
|
||||
t1 = input_text.strip() if strip_whitespace else input_text
|
||||
return (t1, "")
|
||||
|
||||
is_forward = "forward" in str(split_direction).lower()
|
||||
|
||||
if use_regex:
|
||||
try:
|
||||
if is_forward:
|
||||
match = re.search(delimiter, input_text)
|
||||
if match:
|
||||
start, end = match.span()
|
||||
t1 = input_text[:start]
|
||||
t2 = input_text[end:]
|
||||
else:
|
||||
t1, t2 = input_text, ""
|
||||
else:
|
||||
# Backward: find all matches and take the last match
|
||||
matches = list(re.finditer(delimiter, input_text))
|
||||
if matches:
|
||||
last_match = matches[-1]
|
||||
start, end = last_match.span()
|
||||
t1 = input_text[:start]
|
||||
t2 = input_text[end:]
|
||||
else:
|
||||
t1, t2 = input_text, ""
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes Text Split] Regex error: {e}. Falling back to literal string split.")
|
||||
if is_forward:
|
||||
parts = input_text.split(delimiter, 1)
|
||||
else:
|
||||
parts = input_text.rsplit(delimiter, 1)
|
||||
t1 = parts[0]
|
||||
t2 = parts[1] if len(parts) > 1 else ""
|
||||
else:
|
||||
# Literal string sequence split
|
||||
if is_forward:
|
||||
parts = input_text.split(delimiter, 1)
|
||||
else:
|
||||
parts = input_text.rsplit(delimiter, 1)
|
||||
t1 = parts[0]
|
||||
t2 = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
if strip_whitespace:
|
||||
t1 = t1.strip()
|
||||
t2 = t2.strip()
|
||||
|
||||
return (t1, t2)
|
||||
@@ -0,0 +1,259 @@
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import webbrowser
|
||||
from PIL import Image, ImageDraw
|
||||
from .utils import get_saturnnodes_user_dir, get_env_setting
|
||||
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
ENV_FILE = os.path.join(USER_DIR, ".env")
|
||||
|
||||
def is_tray_icon_enabled():
|
||||
val = get_env_setting("ENABLE_TRAY_ICON", "false").lower()
|
||||
if val not in ["true", "1", "yes"]:
|
||||
return False
|
||||
try:
|
||||
import pystray
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
class TrayIconManager:
|
||||
def __init__(self, pause_manager=None):
|
||||
self.pause_manager = pause_manager
|
||||
self._icon = None
|
||||
self._thread = None
|
||||
self._lock = threading.RLock()
|
||||
self._is_running = False
|
||||
|
||||
def create_icon_image(self, color_hex="#059669", symbol="pause"):
|
||||
size = 64
|
||||
img = Image.new("RGBA", (size, size), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
# Draw rounded rectangle matching toolbar button styling
|
||||
draw.rounded_rectangle(
|
||||
[2, 2, size - 3, size - 3],
|
||||
radius=16,
|
||||
fill=color_hex,
|
||||
outline=(0, 0, 0, 60),
|
||||
width=1
|
||||
)
|
||||
|
||||
if symbol == "play":
|
||||
# Play triangle icon (white)
|
||||
draw.polygon([(24, 18), (24, 46), (46, 32)], fill="white")
|
||||
else:
|
||||
# Pause two-bar icon (white)
|
||||
draw.rectangle([20, 18, 28, 46], fill="white")
|
||||
draw.rectangle([36, 18, 44, 46], fill="white")
|
||||
|
||||
return img
|
||||
|
||||
def get_current_colors(self):
|
||||
unpaused_color = get_env_setting("PAUSE_BUTTON_UNPAUSED_COLOR", "#059669")
|
||||
paused_color = get_env_setting("PAUSE_BUTTON_PAUSED_COLOR", "#ea580c")
|
||||
return unpaused_color, paused_color
|
||||
|
||||
def get_status_info(self):
|
||||
if not self.pause_manager:
|
||||
return "Running", "#059669", "pause", "ComfyUI SaturnNodes: Running"
|
||||
|
||||
unpaused_color, paused_color = self.get_current_colors()
|
||||
is_paused = self.pause_manager.paused
|
||||
is_waiting = self.pause_manager.is_waiting
|
||||
mode = self.pause_manager.mode
|
||||
mode_label = "Instant" if mode == "instantly" else "Finish"
|
||||
|
||||
if is_paused:
|
||||
if is_waiting:
|
||||
status_text = "Paused (Waiting)"
|
||||
color = paused_color
|
||||
symbol = "play"
|
||||
title = f"ComfyUI SaturnNodes: Paused ({mode_label}) - Click to Continue"
|
||||
else:
|
||||
status_text = f"Pausing ({mode_label})..."
|
||||
color = paused_color
|
||||
symbol = "pause"
|
||||
title = f"ComfyUI SaturnNodes: Pausing ({mode_label})... (Waiting to complete)"
|
||||
else:
|
||||
status_text = f"Running [Next Pause: {mode_label}]"
|
||||
color = unpaused_color
|
||||
symbol = "pause"
|
||||
title = f"ComfyUI SaturnNodes: Running (Mode: {mode_label})"
|
||||
|
||||
return status_text, color, symbol, title
|
||||
|
||||
def build_menu(self):
|
||||
try:
|
||||
import pystray
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
status_text, _, _, _ = self.get_status_info()
|
||||
is_paused = bool(self.pause_manager and self.pause_manager.paused)
|
||||
is_waiting = bool(self.pause_manager and self.pause_manager.is_waiting)
|
||||
current_mode = self.pause_manager.mode if self.pause_manager else "after_finish"
|
||||
|
||||
def on_toggle_click(icon, item):
|
||||
if not self.pause_manager:
|
||||
return
|
||||
if self.pause_manager.paused:
|
||||
self.pause_manager.set_pause(False)
|
||||
else:
|
||||
self.pause_manager.set_pause(True)
|
||||
|
||||
def on_pause_finish_click(icon, item):
|
||||
if self.pause_manager:
|
||||
self.pause_manager.set_pause(True, mode="after_finish")
|
||||
|
||||
def on_pause_instant_click(icon, item):
|
||||
if self.pause_manager:
|
||||
self.pause_manager.set_pause(True, mode="instantly")
|
||||
|
||||
def on_continue_click(icon, item):
|
||||
if self.pause_manager:
|
||||
self.pause_manager.set_pause(False)
|
||||
|
||||
def on_set_mode_finish(icon, item):
|
||||
if self.pause_manager:
|
||||
self.pause_manager.set_mode("after_finish")
|
||||
|
||||
def on_set_mode_instant(icon, item):
|
||||
if self.pause_manager:
|
||||
self.pause_manager.set_mode("instantly")
|
||||
|
||||
def on_open_browser(icon, item):
|
||||
try:
|
||||
server_url = "http://127.0.0.1:8188"
|
||||
from server import PromptServer
|
||||
if hasattr(PromptServer, "instance") and hasattr(PromptServer.instance, "address"):
|
||||
addr = PromptServer.instance.address
|
||||
if addr and len(addr) >= 2:
|
||||
host = addr[0] or "127.0.0.1"
|
||||
port = addr[1]
|
||||
server_url = f"http://{host}:{port}"
|
||||
webbrowser.open(server_url)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error opening browser: {e}")
|
||||
|
||||
def on_disable_tray(icon, item):
|
||||
self.disable_and_stop()
|
||||
|
||||
if is_paused:
|
||||
mode_desc = "Instant" if current_mode == "instantly" else "After Current"
|
||||
status_display = f"Paused ({mode_desc})"
|
||||
if is_waiting:
|
||||
status_display = f"Pausing ({mode_desc})..."
|
||||
|
||||
menu_items = [
|
||||
pystray.MenuItem(f"● Status: {status_display}", None, enabled=False),
|
||||
pystray.MenuItem("▶ Resume Queue", on_continue_click, default=True),
|
||||
]
|
||||
else:
|
||||
menu_items = [
|
||||
pystray.MenuItem("● Status: Running", None, enabled=False),
|
||||
pystray.MenuItem("⏸ Pause (Instant)", on_pause_instant_click, default=False),
|
||||
pystray.MenuItem("⏳ Pause (After Current)", on_pause_finish_click, default=False),
|
||||
]
|
||||
|
||||
return pystray.Menu(*menu_items)
|
||||
|
||||
def start(self):
|
||||
with self._lock:
|
||||
if self._is_running or self._icon is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
import pystray
|
||||
except ImportError:
|
||||
print("[SaturnNodes] System tray icon disabled: 'pystray' is not installed. Install pystray ('pip install pystray') to enable system tray support.")
|
||||
return
|
||||
|
||||
_, color, symbol, title = self.get_status_info()
|
||||
icon_img = self.create_icon_image(color, symbol)
|
||||
menu = self.build_menu()
|
||||
|
||||
self._icon = pystray.Icon(
|
||||
name="ComfyUI-SaturnNodes",
|
||||
icon=icon_img,
|
||||
title=title,
|
||||
menu=menu
|
||||
)
|
||||
|
||||
def _run():
|
||||
try:
|
||||
self._is_running = True
|
||||
self._icon.run()
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] System tray icon error: {e}")
|
||||
finally:
|
||||
self._is_running = False
|
||||
self._icon = None
|
||||
|
||||
self._thread = threading.Thread(target=_run, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def stop(self):
|
||||
with self._lock:
|
||||
if self._icon is not None:
|
||||
try:
|
||||
self._icon.stop()
|
||||
except Exception:
|
||||
pass
|
||||
self._icon = None
|
||||
self._is_running = False
|
||||
print("[SaturnNodes] System tray icon stopped.")
|
||||
|
||||
def update_status(self):
|
||||
with self._lock:
|
||||
if not self._is_running or self._icon is None:
|
||||
return
|
||||
|
||||
try:
|
||||
_, color, symbol, title = self.get_status_info()
|
||||
self._icon.icon = self.create_icon_image(color, symbol)
|
||||
self._icon.title = title
|
||||
self._icon.menu = self.build_menu()
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error updating system tray status: {e}")
|
||||
|
||||
def set_enabled(self, enabled):
|
||||
if enabled:
|
||||
if not self._is_running:
|
||||
self.start()
|
||||
else:
|
||||
self.update_status()
|
||||
else:
|
||||
if self._is_running:
|
||||
self.stop()
|
||||
|
||||
def disable_and_stop(self):
|
||||
# Update .env setting and stop tray
|
||||
lines = []
|
||||
if os.path.exists(ENV_FILE):
|
||||
try:
|
||||
with open(ENV_FILE, "r", encoding="utf-8") as f:
|
||||
lines = f.readlines()
|
||||
except Exception:
|
||||
lines = []
|
||||
|
||||
new_lines = []
|
||||
has_tray = False
|
||||
for line in lines:
|
||||
if line.strip().startswith("ENABLE_TRAY_ICON="):
|
||||
new_lines.append("ENABLE_TRAY_ICON=false\n")
|
||||
has_tray = True
|
||||
else:
|
||||
new_lines.append(line)
|
||||
|
||||
if not has_tray:
|
||||
new_lines.append("ENABLE_TRAY_ICON=false\n")
|
||||
|
||||
try:
|
||||
with open(ENV_FILE, "w", encoding="utf-8") as f:
|
||||
f.writelines(new_lines)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.stop()
|
||||
@@ -1,80 +0,0 @@
|
||||
import os
|
||||
import re
|
||||
import folder_paths
|
||||
|
||||
UTILS_CATEGORY = "🍃 FlowControl/Utils"
|
||||
|
||||
def parse_pretty_name(filepath):
|
||||
base = os.path.splitext(os.path.basename(filepath))[0]
|
||||
parts = base.split('_')
|
||||
if len(parts) >= 2:
|
||||
name_part = parts[1]
|
||||
name_part = re.sub(r'(?<!^)(?=[A-Z])', ' ', name_part)
|
||||
name_part = name_part.replace('-', ' ')
|
||||
|
||||
words = name_part.split()
|
||||
formatted_words = []
|
||||
for word in words:
|
||||
if word.upper() in ["NSFW", "LORA", "V1", "V2", "V3", "V4", "FP16", "HM"]:
|
||||
formatted_words.append(word.upper())
|
||||
else:
|
||||
formatted_words.append(word.capitalize())
|
||||
return " ".join(formatted_words)
|
||||
return base.replace('-', ' ')
|
||||
|
||||
def get_pretty_names_for_folder(folder):
|
||||
all_loras = folder_paths.get_filename_list("loras")
|
||||
raw_filter = os.path.normpath(folder).lower().replace("\\", "/")
|
||||
raw_filter = raw_filter.rstrip("/* ")
|
||||
|
||||
pretty_names = set()
|
||||
for lora in all_loras:
|
||||
lora_path = folder_paths.get_full_path("loras", lora)
|
||||
if not lora_path:
|
||||
continue
|
||||
|
||||
abs_lora = os.path.normpath(lora_path).lower().replace("\\", "/")
|
||||
rel_lora = lora.lower().replace("\\", "/")
|
||||
|
||||
if raw_filter:
|
||||
if not abs_lora.startswith(raw_filter) and not rel_lora.startswith(raw_filter):
|
||||
continue
|
||||
|
||||
pretty = parse_pretty_name(lora)
|
||||
if pretty and pretty not in ["[ NONE ]", "[ RANDOM ]"]:
|
||||
pretty_names.add(pretty)
|
||||
|
||||
return sorted(list(pretty_names), key=len, reverse=True)
|
||||
|
||||
class UndoPlaceholder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {"forceInput": True}),
|
||||
"lora_folder": ("STRING", {"default": ""}),
|
||||
"placeholder": ("STRING", {"default": "%celeb%"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("prompt",)
|
||||
FUNCTION = "undo_placeholder"
|
||||
CATEGORY = UTILS_CATEGORY
|
||||
DESCRIPTION = "Scans prompt string for names matching LoRA files in the given folder and replaces them back with a placeholder token (e.g. %celeb%)."
|
||||
|
||||
def undo_placeholder(self, prompt, lora_folder, placeholder):
|
||||
if not prompt or not lora_folder:
|
||||
return (prompt,)
|
||||
|
||||
pretty_names = get_pretty_names_for_folder(lora_folder)
|
||||
if not pretty_names:
|
||||
return (prompt,)
|
||||
|
||||
modified_prompt = prompt
|
||||
for name in pretty_names:
|
||||
escaped_name = re.escape(name)
|
||||
pattern = re.compile(rf'\b{escaped_name}\b', re.IGNORECASE)
|
||||
modified_prompt = pattern.sub(placeholder, modified_prompt)
|
||||
|
||||
return (modified_prompt,)
|
||||
+452
@@ -0,0 +1,452 @@
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import threading
|
||||
from urllib.parse import urlparse
|
||||
import folder_paths
|
||||
|
||||
def get_saturnnodes_user_dir():
|
||||
"""
|
||||
Returns the centralized user data directory: ComfyUI/user/default/SaturnNodes/
|
||||
"""
|
||||
base_user = None
|
||||
try:
|
||||
if hasattr(folder_paths, "get_user_directory"):
|
||||
user_base = folder_paths.get_user_directory()
|
||||
if user_base:
|
||||
norm_base = os.path.normpath(user_base)
|
||||
if os.path.basename(norm_base).lower() == "user":
|
||||
base_user = os.path.join(norm_base, "default")
|
||||
else:
|
||||
base_user = norm_base
|
||||
except Exception:
|
||||
base_user = None
|
||||
|
||||
if not base_user:
|
||||
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||||
base_user = os.path.join(root_dir, "user", "default")
|
||||
|
||||
saturn_dir = os.path.join(base_user, "SaturnNodes")
|
||||
os.makedirs(saturn_dir, exist_ok=True)
|
||||
return saturn_dir
|
||||
|
||||
_SESSION_CSRF_TOKEN = secrets.token_hex(32)
|
||||
|
||||
def get_csrf_token():
|
||||
"""Returns the ephemeral cryptographically random session CSRF token generated on server boot."""
|
||||
return _SESSION_CSRF_TOKEN
|
||||
|
||||
def is_same_origin_or_direct(request):
|
||||
"""
|
||||
Validates that a request originates from the local ComfyUI interface and not from a visited web page (CSRF defense).
|
||||
"""
|
||||
if request is None:
|
||||
return True
|
||||
|
||||
# 1. Reject cross-site requests via Fetch Metadata (modern browsers always send this; cannot be spoofed by JS)
|
||||
sec_fetch_site = getattr(request, "headers", {}).get("Sec-Fetch-Site")
|
||||
if sec_fetch_site and sec_fetch_site.lower() == "cross-site":
|
||||
return False
|
||||
|
||||
# 2. Check Origin header if present
|
||||
origin = getattr(request, "headers", {}).get("Origin")
|
||||
if origin:
|
||||
try:
|
||||
parsed = urlparse(origin)
|
||||
host = parsed.hostname
|
||||
if host not in ["127.0.0.1", "localhost", "::1", "0.0.0.0", "testclient"]:
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# 3. Check Referer header if present
|
||||
referer = getattr(request, "headers", {}).get("Referer")
|
||||
if referer:
|
||||
try:
|
||||
parsed_ref = urlparse(referer)
|
||||
ref_host = parsed_ref.hostname
|
||||
if ref_host not in ["127.0.0.1", "localhost", "::1", "0.0.0.0", "testclient", "::ffff:127.0.0.1"]:
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_authenticated_local_request(request):
|
||||
"""
|
||||
Guarantees user-initiated actions and protects against CSRF / drive-by requests from open browser tabs:
|
||||
1. Originates strictly from loopback IP (is_local_request)
|
||||
2. Not a cross-site drive-by request (is_same_origin_or_direct)
|
||||
3. Contains a valid X-SaturnNodes-CSRF-Token header matching the session secret.
|
||||
"""
|
||||
if not is_local_request(request):
|
||||
return False
|
||||
if not is_same_origin_or_direct(request):
|
||||
return False
|
||||
|
||||
headers = getattr(request, "headers", {})
|
||||
token = headers.get("X-SaturnNodes-CSRF-Token")
|
||||
if not token or not secrets.compare_digest(token, _SESSION_CSRF_TOKEN):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_local_request(request):
|
||||
"""
|
||||
Returns True if the HTTP request originates strictly from loopback (localhost / 127.0.0.1 / ::1 / ::ffff:127.0.0.1).
|
||||
Used to protect sensitive endpoints (power control, system settings, file browsing).
|
||||
"""
|
||||
if request is None:
|
||||
return True
|
||||
remote = getattr(request, "remote", None)
|
||||
if not remote:
|
||||
transport = getattr(request, "transport", None)
|
||||
if transport:
|
||||
peername = transport.get_extra_info("peername")
|
||||
if peername and isinstance(peername, tuple):
|
||||
remote = peername[0]
|
||||
if not remote:
|
||||
return False
|
||||
clean_remote = str(remote).strip("[]")
|
||||
if ":" in clean_remote and clean_remote.count(":") == 1:
|
||||
clean_remote = clean_remote.split(":")[0]
|
||||
local_hosts = {"127.0.0.1", "::1", "localhost", "0.0.0.0", "testclient", "::ffff:127.0.0.1"}
|
||||
if clean_remote in local_hosts or remote in local_hosts:
|
||||
return True
|
||||
try:
|
||||
import ipaddress
|
||||
ip = ipaddress.ip_address(clean_remote)
|
||||
return ip.is_loopback
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def get_allowed_image_directories():
|
||||
"""
|
||||
Returns approved base directories for image loader operations: strictly ComfyUI input, output, and temp dirs.
|
||||
"""
|
||||
allowed = []
|
||||
try:
|
||||
inp = folder_paths.get_input_directory()
|
||||
if inp:
|
||||
allowed.append(os.path.abspath(os.path.realpath(inp)))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
out = folder_paths.get_output_directory()
|
||||
if out:
|
||||
allowed.append(os.path.abspath(os.path.realpath(out)))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
tmp = folder_paths.get_temp_directory()
|
||||
if tmp:
|
||||
allowed.append(os.path.abspath(os.path.realpath(tmp)))
|
||||
except Exception:
|
||||
pass
|
||||
return allowed
|
||||
|
||||
def is_safe_path(target_path, allowed_bases=None):
|
||||
"""
|
||||
Confines file access strictly within approved base directories.
|
||||
Prevents path traversal, directory escape, and arbitrary system file reads.
|
||||
"""
|
||||
if not target_path:
|
||||
return False
|
||||
if allowed_bases is None:
|
||||
allowed_bases = get_allowed_image_directories()
|
||||
try:
|
||||
abs_target = os.path.abspath(os.path.realpath(target_path))
|
||||
for base in allowed_bases:
|
||||
abs_base = os.path.abspath(os.path.realpath(base))
|
||||
if os.path.commonpath([abs_base, abs_target]) == abs_base:
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
return False
|
||||
|
||||
def sanitize_image_loader_folder(folder_input, default_to_output=True):
|
||||
"""
|
||||
Resolves folder input strictly within ComfyUI input/output/temp directories.
|
||||
Rejects directory traversal ('..') and unapproved absolute paths, strictly verifying confinement via commonpath.
|
||||
"""
|
||||
allowed_bases = get_allowed_image_directories()
|
||||
input_dir = folder_paths.get_input_directory() if hasattr(folder_paths, "get_input_directory") else "."
|
||||
output_dir = folder_paths.get_output_directory() if hasattr(folder_paths, "get_output_directory") else "."
|
||||
default_dir = os.path.abspath(os.path.realpath(output_dir if default_to_output else input_dir))
|
||||
|
||||
if not folder_input or not str(folder_input).strip():
|
||||
return default_dir
|
||||
|
||||
clean = str(folder_input).strip().rstrip("\\/")
|
||||
|
||||
# 1. Reject directory traversal sequences ('..')
|
||||
norm_slashes = clean.replace("\\", "/")
|
||||
if any(p == ".." for p in norm_slashes.split("/")):
|
||||
print(f"[SaturnNodes Security] Blocked directory traversal in image folder path '{clean}'. Confining to ComfyUI directories.")
|
||||
return default_dir
|
||||
|
||||
# 2. Reject absolute paths (drive letters C:\, POSIX root /, UNC \\server)
|
||||
is_absolute = (
|
||||
os.path.isabs(clean)
|
||||
or bool(re.match(r'^[A-Za-z]:', clean))
|
||||
or clean.startswith(("//", "\\\\"))
|
||||
)
|
||||
if is_absolute:
|
||||
if not is_safe_path(clean, allowed_bases):
|
||||
print(f"[SaturnNodes Security] Blocked absolute image folder path '{clean}'. Confining to ComfyUI directories.")
|
||||
return default_dir
|
||||
return os.path.abspath(os.path.realpath(clean))
|
||||
|
||||
# 3. Handle relative 'input' or 'output' prefixes
|
||||
if clean.lower() == "input" or clean.lower().startswith("input/") or clean.lower().startswith("input\\"):
|
||||
sub = clean[5:].lstrip("\\/")
|
||||
resolved = os.path.join(input_dir, sub) if sub else input_dir
|
||||
elif clean.lower() == "output" or clean.lower().startswith("output/") or clean.lower().startswith("output\\"):
|
||||
sub = clean[6:].lstrip("\\/")
|
||||
resolved = os.path.join(output_dir, sub) if sub else output_dir
|
||||
else:
|
||||
cand_out = os.path.join(output_dir, clean)
|
||||
cand_inp = os.path.join(input_dir, clean)
|
||||
if os.path.exists(cand_inp) and not os.path.exists(cand_out):
|
||||
resolved = cand_inp
|
||||
else:
|
||||
resolved = cand_out
|
||||
|
||||
if is_safe_path(resolved, allowed_bases):
|
||||
return os.path.abspath(os.path.realpath(resolved))
|
||||
|
||||
return default_dir
|
||||
|
||||
_ENV_CACHE_LOCK = threading.RLock()
|
||||
_ENV_CACHE = None
|
||||
_ENV_CACHE_MTIME = None
|
||||
|
||||
def invalidate_env_settings_cache():
|
||||
"""Force the next get_env_setting() call to re-read the .env file."""
|
||||
global _ENV_CACHE, _ENV_CACHE_MTIME
|
||||
with _ENV_CACHE_LOCK:
|
||||
_ENV_CACHE = None
|
||||
_ENV_CACHE_MTIME = None
|
||||
|
||||
def get_env_setting(key, default_val):
|
||||
"""
|
||||
Reads a key setting from ComfyUI/user/default/SaturnNodes/.env safely.
|
||||
Parsed values are cached until the file mtime changes.
|
||||
"""
|
||||
global _ENV_CACHE, _ENV_CACHE_MTIME
|
||||
user_dir = get_saturnnodes_user_dir()
|
||||
env_file = os.path.join(user_dir, ".env")
|
||||
|
||||
mtime = None
|
||||
try:
|
||||
if os.path.exists(env_file):
|
||||
mtime = os.stat(env_file).st_mtime_ns
|
||||
except OSError:
|
||||
mtime = None
|
||||
|
||||
with _ENV_CACHE_LOCK:
|
||||
if _ENV_CACHE is None or mtime != _ENV_CACHE_MTIME:
|
||||
parsed = {}
|
||||
if mtime is not None:
|
||||
try:
|
||||
with open(env_file, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
stripped = line.strip()
|
||||
if not stripped or stripped.startswith("#") or "=" not in stripped:
|
||||
continue
|
||||
k, v = stripped.split("=", 1)
|
||||
parsed[k.strip()] = v.strip()
|
||||
except Exception:
|
||||
parsed = {}
|
||||
_ENV_CACHE = parsed
|
||||
_ENV_CACHE_MTIME = mtime
|
||||
|
||||
return _ENV_CACHE.get(key, default_val)
|
||||
|
||||
def parse_pretty_name(filepath):
|
||||
if not filepath or filepath in ["[ NONE ]", "[ RANDOM ]"]:
|
||||
return ""
|
||||
base = os.path.splitext(os.path.basename(filepath))[0]
|
||||
parts = base.split('_')
|
||||
if len(parts) >= 2:
|
||||
name_part = parts[1]
|
||||
name_part = re.sub(r'(?<=[a-z])(?=[A-Z])', ' ', name_part)
|
||||
name_part = name_part.replace('-', ' ')
|
||||
|
||||
words = name_part.split()
|
||||
formatted_words = []
|
||||
for word in words:
|
||||
if word.upper() in ["NSFW", "LORA", "V1", "V2", "V3", "V4", "FP16", "HM"]:
|
||||
formatted_words.append(word.upper())
|
||||
else:
|
||||
formatted_words.append(word.capitalize())
|
||||
return " ".join(formatted_words)
|
||||
return base.replace('-', ' ')
|
||||
|
||||
def parse_pretty_name_with_version(filepath):
|
||||
if not filepath or filepath in ["[ NONE ]", "[ RANDOM ]"]:
|
||||
return ""
|
||||
base = os.path.splitext(os.path.basename(filepath))[0]
|
||||
parts = base.split('_')
|
||||
pretty_name = parse_pretty_name(filepath)
|
||||
|
||||
if len(parts) >= 3:
|
||||
version_candidate = parts[2].strip()
|
||||
if re.match(r'^v[0-9]+(\.[0-9]+)?$', version_candidate, re.IGNORECASE):
|
||||
return f"{pretty_name} {version_candidate.upper()}"
|
||||
return pretty_name
|
||||
|
||||
def get_allowed_folder_watch_directories():
|
||||
"""
|
||||
Returns approved base directories for folder watching and recent output scanning:
|
||||
strictly ComfyUI input and output directories.
|
||||
"""
|
||||
allowed = []
|
||||
try:
|
||||
inp = folder_paths.get_input_directory()
|
||||
if inp:
|
||||
allowed.append(os.path.abspath(os.path.realpath(inp)))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
out = folder_paths.get_output_directory()
|
||||
if out:
|
||||
allowed.append(os.path.abspath(os.path.realpath(out)))
|
||||
except Exception:
|
||||
pass
|
||||
return allowed
|
||||
|
||||
def sanitize_folder_path(folder_input, default_dir=None):
|
||||
"""
|
||||
Confines folder paths strictly within ComfyUI's input and output directories.
|
||||
Rejects absolute paths and '..' directory traversal sequences.
|
||||
Resolves canonical realpath and verifies commonpath containment.
|
||||
"""
|
||||
allowed_bases = get_allowed_folder_watch_directories()
|
||||
input_dir = folder_paths.get_input_directory() if hasattr(folder_paths, "get_input_directory") else None
|
||||
safe_fallback = os.path.join(input_dir, "watch") if input_dir else "watch"
|
||||
if default_dir:
|
||||
if os.path.isabs(default_dir):
|
||||
try:
|
||||
real_def = os.path.abspath(os.path.realpath(default_dir))
|
||||
for base in allowed_bases:
|
||||
if os.path.commonpath([base, real_def]) == base:
|
||||
safe_fallback = real_def
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
safe_fallback = default_dir
|
||||
|
||||
if folder_input is None:
|
||||
folder_input = ""
|
||||
clean_folder = str(folder_input).strip()
|
||||
|
||||
# Strip any trailing wildcards (e.g. "subfolder\*" -> "subfolder", "subfolder/*" -> "subfolder")
|
||||
clean_folder = re.sub(r'[\*\?]+$', '', clean_folder).rstrip("\\/")
|
||||
|
||||
if not clean_folder:
|
||||
clean_folder = safe_fallback
|
||||
|
||||
# 1. Reject absolute paths (drive letters C:\, POSIX root /, UNC \\server)
|
||||
is_absolute = (
|
||||
os.path.isabs(clean_folder)
|
||||
or bool(re.match(r'^[A-Za-z]:', clean_folder))
|
||||
or clean_folder.startswith(("//", "\\\\"))
|
||||
)
|
||||
|
||||
# 2. Reject directory traversal sequences ('..')
|
||||
normalized_slashes = clean_folder.replace("\\", "/")
|
||||
parts = normalized_slashes.split("/")
|
||||
has_traversal = any(p == ".." for p in parts)
|
||||
|
||||
if is_absolute or has_traversal:
|
||||
is_safe_default = False
|
||||
if is_absolute and safe_fallback and clean_folder == safe_fallback:
|
||||
try:
|
||||
real_clean = os.path.abspath(os.path.realpath(clean_folder))
|
||||
for base in allowed_bases:
|
||||
if os.path.commonpath([base, real_clean]) == base:
|
||||
is_safe_default = True
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not is_safe_default:
|
||||
print(f"[SaturnNodes Security] Blocked unsafe folder path '{clean_folder}' (absolute={is_absolute}, traversal={has_traversal}). Confining to ComfyUI input/output directories.")
|
||||
clean_folder = safe_fallback
|
||||
|
||||
# 3. If clean_folder is an already-confined absolute path (e.g. from internal safe_fallback)
|
||||
if os.path.isabs(clean_folder):
|
||||
try:
|
||||
real_cand = os.path.abspath(os.path.realpath(clean_folder))
|
||||
for base in allowed_bases:
|
||||
if os.path.commonpath([base, real_cand]) == base:
|
||||
return real_cand
|
||||
except Exception:
|
||||
pass
|
||||
return os.path.abspath(os.path.realpath(safe_fallback))
|
||||
|
||||
# 4. For relative paths, evaluate candidates strictly within allowed_bases
|
||||
candidates = []
|
||||
for base in allowed_bases:
|
||||
target = os.path.abspath(os.path.realpath(os.path.join(base, clean_folder)))
|
||||
try:
|
||||
if os.path.commonpath([base, target]) == base:
|
||||
candidates.append((target, os.path.exists(target)))
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# Return existing candidate if present on disk
|
||||
for target, exists in candidates:
|
||||
if exists:
|
||||
return target
|
||||
|
||||
# Case-insensitive resolution within candidate parent folders
|
||||
for target, _ in candidates:
|
||||
parent_dir = os.path.dirname(target)
|
||||
base_name = os.path.basename(target)
|
||||
if os.path.exists(parent_dir):
|
||||
try:
|
||||
for entry in os.listdir(parent_dir):
|
||||
if entry.lower() == base_name.lower():
|
||||
match_path = os.path.join(parent_dir, entry)
|
||||
real_match = os.path.abspath(os.path.realpath(match_path))
|
||||
for base in allowed_bases:
|
||||
if os.path.commonpath([base, real_match]) == base:
|
||||
return real_match
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if candidates:
|
||||
return candidates[0][0]
|
||||
|
||||
default_input = folder_paths.get_input_directory() if hasattr(folder_paths, "get_input_directory") else "."
|
||||
return os.path.abspath(os.path.realpath(os.path.join(default_input, "watch")))
|
||||
|
||||
def format_lora_output_name(resolved_path, display_name, output_format="Parsed Name", custom_regex=""):
|
||||
if not resolved_path or resolved_path in ["[ NONE ]", "[ RANDOM ]"]:
|
||||
return ""
|
||||
|
||||
if output_format == "Filename":
|
||||
return os.path.basename(resolved_path)
|
||||
elif output_format == "Filename without extension":
|
||||
return os.path.splitext(os.path.basename(resolved_path))[0]
|
||||
elif output_format == "Relative Path":
|
||||
return resolved_path.replace("\\", "/")
|
||||
elif output_format == "Full Path":
|
||||
full_p = folder_paths.get_full_path("loras", resolved_path)
|
||||
return full_p.replace("\\", "/") if full_p else resolved_path
|
||||
elif output_format == "Custom Regex" and custom_regex and custom_regex.strip():
|
||||
try:
|
||||
m = re.search(custom_regex, resolved_path)
|
||||
if m:
|
||||
return m.group(0)
|
||||
except Exception:
|
||||
pass
|
||||
return os.path.splitext(os.path.basename(resolved_path))[0]
|
||||
else: # "Parsed Name" (default)
|
||||
clean_name = display_name
|
||||
if " - " in display_name:
|
||||
clean_name = display_name.split(" - ", 1)[1]
|
||||
return re.sub(r'\s+V\d+(\.\d+)?$', '', clean_name, flags=re.IGNORECASE).strip()
|
||||
+21
-8
@@ -1,14 +1,27 @@
|
||||
[project]
|
||||
name = "comfyui-flowcontrol"
|
||||
name = "saturnnodes"
|
||||
description = "Essential workflow controls, mid-generation pausing, real-time queue persistence, folder watcher, and visual model loaders for ComfyUI."
|
||||
version = "1.0.0"
|
||||
license = "MIT"
|
||||
dependencies = ["torch", "pillow", "numpy", "piexif"]
|
||||
version = "2.6.3"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
dependencies = ["pillow", "numpy", "piexif"]
|
||||
classifiers = [
|
||||
"Operating System :: OS Independent",
|
||||
"Programming Language :: Python :: 3",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
tray = ["pystray"]
|
||||
stereogram = ["onnxruntime"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/KOFiblto/ComfyUI-FlowControl"
|
||||
Repository = "https://github.com/KOFiblto/ComfyUI-SaturnNodes"
|
||||
"Bug Tracker" = "https://github.com/KOFiblto/ComfyUI-SaturnNodes/issues"
|
||||
Documentation = "https://github.com/KOFiblto/ComfyUI-SaturnNodes#readme"
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "KOFiblto"
|
||||
DisplayName = "ComfyUI-FlowControl"
|
||||
Icon = ""
|
||||
PublisherId = "kofiblto"
|
||||
DisplayName = "ComfyUI-SaturnNodes"
|
||||
Icon = "https://raw.githubusercontent.com/KOFiblto/ComfyUI-SaturnNodes/main/web/logo.png"
|
||||
requires-comfyui = ">=0.3.0"
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
Pillow
|
||||
piexif
|
||||
numpy
|
||||
# pystray # Optional: install manually if desktop system tray icon integration is enabled
|
||||
# onnxruntime # Optional: install manually for built-in Depth Anything V2 estimation in 3D stereogram nodes
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Ensure SaturnNodes root is at index 0 in sys.path
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if ROOT in sys.path:
|
||||
sys.path.remove(ROOT)
|
||||
sys.path.insert(0, ROOT)
|
||||
@@ -0,0 +1,539 @@
|
||||
// Minimal DOM used by the standalone frontend regression tests.
|
||||
// It implements only the APIs the SaturnNodes extensions touch.
|
||||
|
||||
function splitTopLevel(selector, delimiter = ",") {
|
||||
const parts = [];
|
||||
let depth = 0;
|
||||
let quote = null;
|
||||
let current = "";
|
||||
for (const ch of String(selector)) {
|
||||
if (quote) {
|
||||
current += ch;
|
||||
if (ch === quote) quote = null;
|
||||
continue;
|
||||
}
|
||||
if (ch === '"' || ch === "'") {
|
||||
quote = ch;
|
||||
current += ch;
|
||||
continue;
|
||||
}
|
||||
if (ch === "[" || ch === "(") depth++;
|
||||
if (ch === "]" || ch === ")") depth--;
|
||||
if (ch === delimiter && depth === 0) {
|
||||
parts.push(current.trim());
|
||||
current = "";
|
||||
continue;
|
||||
}
|
||||
current += ch;
|
||||
}
|
||||
if (current.trim()) parts.push(current.trim());
|
||||
return parts.filter(Boolean);
|
||||
}
|
||||
|
||||
function parseAttributeTest(raw) {
|
||||
const match = raw.match(
|
||||
/^\[\s*([\w-]+)\s*(?:([*^$~|]?=)\s*(?:"([^"]*)"|'([^']*)'|([^\]\s]+))\s*)?(?:([iIsS]))?\s*\]$/
|
||||
);
|
||||
if (!match) return null;
|
||||
return {
|
||||
name: match[1],
|
||||
operator: match[2] || null,
|
||||
value: match[3] ?? match[4] ?? match[5] ?? "",
|
||||
insensitive: (match[6] || "").toLowerCase() === "i",
|
||||
};
|
||||
}
|
||||
|
||||
function matchAttribute(el, test) {
|
||||
if (!el.hasAttribute(test.name)) return false;
|
||||
if (!test.operator) return true;
|
||||
let actual = String(el.getAttribute(test.name));
|
||||
let expected = test.value;
|
||||
if (test.insensitive) {
|
||||
actual = actual.toLowerCase();
|
||||
expected = expected.toLowerCase();
|
||||
}
|
||||
switch (test.operator) {
|
||||
case "=": return actual === expected;
|
||||
case "*=": return actual.includes(expected);
|
||||
case "^=": return actual.startsWith(expected);
|
||||
case "$=": return actual.endsWith(expected);
|
||||
case "~=": return actual.split(/\s+/).includes(expected);
|
||||
case "|=": return actual === expected || actual.startsWith(`${expected}-`);
|
||||
default: return false;
|
||||
}
|
||||
}
|
||||
|
||||
function matchCompound(el, compound) {
|
||||
if (!el || el.nodeType !== 1) return false;
|
||||
let rest = compound.trim();
|
||||
if (!rest) return false;
|
||||
|
||||
const tagMatch = rest.match(/^([a-zA-Z][\w-]*|\*)/);
|
||||
if (tagMatch) {
|
||||
const tag = tagMatch[1];
|
||||
if (tag !== "*" && el.tagName !== tag.toUpperCase()) return false;
|
||||
rest = rest.slice(tag.length);
|
||||
}
|
||||
|
||||
while (rest.length) {
|
||||
if (rest.startsWith("#")) {
|
||||
const m = rest.match(/^#([\w-]+)/);
|
||||
if (!m) return false;
|
||||
if (el.id !== m[1]) return false;
|
||||
rest = rest.slice(m[0].length);
|
||||
continue;
|
||||
}
|
||||
if (rest.startsWith(".")) {
|
||||
const m = rest.match(/^\.([\w-]+)/);
|
||||
if (!m) return false;
|
||||
if (!el.classList.contains(m[1])) return false;
|
||||
rest = rest.slice(m[0].length);
|
||||
continue;
|
||||
}
|
||||
if (rest.startsWith("[")) {
|
||||
let depth = 0;
|
||||
let end = -1;
|
||||
for (let i = 0; i < rest.length; i++) {
|
||||
if (rest[i] === "[") depth++;
|
||||
if (rest[i] === "]") {
|
||||
depth--;
|
||||
if (depth === 0) { end = i; break; }
|
||||
}
|
||||
}
|
||||
if (end < 0) return false;
|
||||
const test = parseAttributeTest(rest.slice(0, end + 1));
|
||||
if (!test || !matchAttribute(el, test)) return false;
|
||||
rest = rest.slice(end + 1);
|
||||
continue;
|
||||
}
|
||||
if (rest.startsWith(":")) {
|
||||
const m = rest.match(/^:([\w-]+)(?:\([^)]*\))?/);
|
||||
if (!m) return false;
|
||||
// :hover / :focus are never true in these tests. Other pseudos are
|
||||
// treated as non-restrictive so simple selectors still work.
|
||||
if (m[1] === "hover" || m[1] === "focus" || m[1] === "focus-visible") return false;
|
||||
rest = rest.slice(m[0].length);
|
||||
continue;
|
||||
}
|
||||
// Unsupported combinator (e.g. ">") - stop rather than mis-match.
|
||||
break;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
function matchChain(el, tokens, index) {
|
||||
if (!matchCompound(el, tokens[index])) return false;
|
||||
if (index === 0) return true;
|
||||
let ancestor = el.parentElement;
|
||||
while (ancestor) {
|
||||
if (matchChain(ancestor, tokens, index - 1)) return true;
|
||||
ancestor = ancestor.parentElement;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
export function matchesSelector(el, selector) {
|
||||
if (!el || el.nodeType !== 1) return false;
|
||||
return splitTopLevel(selector).some(part => {
|
||||
const tokens = part.trim().split(/\s+/).filter(Boolean);
|
||||
if (!tokens.length) return false;
|
||||
return matchChain(el, tokens, tokens.length - 1);
|
||||
});
|
||||
}
|
||||
|
||||
function makeStyle() {
|
||||
const props = new Map();
|
||||
const style = {
|
||||
setProperty(name, value) {
|
||||
props.set(String(name), String(value));
|
||||
const camel = String(name).replace(/-([a-z])/g, (_, c) => c.toUpperCase());
|
||||
style[camel] = String(value);
|
||||
},
|
||||
removeProperty(name) {
|
||||
props.delete(String(name));
|
||||
const camel = String(name).replace(/-([a-z])/g, (_, c) => c.toUpperCase());
|
||||
delete style[camel];
|
||||
},
|
||||
getPropertyValue(name) {
|
||||
return props.get(String(name)) || "";
|
||||
},
|
||||
cssText: "",
|
||||
};
|
||||
return style;
|
||||
}
|
||||
|
||||
export class FakeElement {
|
||||
constructor(tagName, doc = null) {
|
||||
this.tagName = String(tagName || "div").toUpperCase();
|
||||
this.nodeType = 1;
|
||||
this.ownerDocument = doc;
|
||||
this.attributes = new Map();
|
||||
this.children = [];
|
||||
this.parentNode = null;
|
||||
this._classes = new Set();
|
||||
this._listeners = new Map();
|
||||
this._rect = { left: 0, top: 0, right: 0, bottom: 0, width: 0, height: 0 };
|
||||
this._isRoot = false;
|
||||
this._innerHTML = "";
|
||||
this._text = "";
|
||||
this.dataset = {};
|
||||
this.style = makeStyle();
|
||||
this.id = "";
|
||||
this.value = "";
|
||||
this.offsetParent = null;
|
||||
this.offsetHeight = 0;
|
||||
this.offsetWidth = 0;
|
||||
this.scrollTop = 0;
|
||||
this.scrollHeight = 0;
|
||||
this.clientHeight = 0;
|
||||
this.disabled = false;
|
||||
this.type = "";
|
||||
this.title = "";
|
||||
this.src = "";
|
||||
this.checked = false;
|
||||
}
|
||||
|
||||
get className() {
|
||||
return Array.from(this._classes).join(" ");
|
||||
}
|
||||
|
||||
set className(value) {
|
||||
this._classes = new Set(String(value || "").split(/\s+/).filter(Boolean));
|
||||
}
|
||||
|
||||
get classList() {
|
||||
const el = this;
|
||||
return {
|
||||
add(...names) { names.forEach(n => n && el._classes.add(String(n))); },
|
||||
remove(...names) { names.forEach(n => el._classes.delete(String(n))); },
|
||||
contains(name) { return el._classes.has(String(name)); },
|
||||
toggle(name, force) {
|
||||
const has = el._classes.has(String(name));
|
||||
const next = force === undefined ? !has : !!force;
|
||||
if (next) el._classes.add(String(name));
|
||||
else el._classes.delete(String(name));
|
||||
return next;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
get isConnected() {
|
||||
let node = this;
|
||||
while (node) {
|
||||
if (node._isRoot) return true;
|
||||
node = node.parentNode;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
get parentElement() {
|
||||
return this.parentNode && this.parentNode.nodeType === 1 ? this.parentNode : null;
|
||||
}
|
||||
|
||||
get innerHTML() {
|
||||
return this._innerHTML;
|
||||
}
|
||||
|
||||
set innerHTML(value) {
|
||||
this._innerHTML = String(value ?? "");
|
||||
this.children.forEach(child => { child.parentNode = null; });
|
||||
this.children = [];
|
||||
}
|
||||
|
||||
get textContent() {
|
||||
return this._text;
|
||||
}
|
||||
|
||||
set textContent(value) {
|
||||
this._text = String(value ?? "");
|
||||
}
|
||||
|
||||
getBoundingClientRect() {
|
||||
return { ...this._rect };
|
||||
}
|
||||
|
||||
setRect(rect) {
|
||||
this._rect = { ...this._rect, ...rect };
|
||||
this.offsetHeight = this._rect.height;
|
||||
this.offsetWidth = this._rect.width;
|
||||
return this;
|
||||
}
|
||||
|
||||
setAttribute(name, value) {
|
||||
const key = String(name);
|
||||
const stringValue = String(value);
|
||||
this.attributes.set(key, stringValue);
|
||||
if (key === "id") this.id = stringValue;
|
||||
if (key === "class") this.className = stringValue;
|
||||
if (key === "value") this.value = stringValue;
|
||||
if (key === "src") this.src = stringValue;
|
||||
if (key.startsWith("data-")) {
|
||||
const camel = key.slice(5).replace(/-([a-z])/g, (_, c) => c.toUpperCase());
|
||||
this.dataset[camel] = stringValue;
|
||||
}
|
||||
}
|
||||
|
||||
getAttribute(name) {
|
||||
const key = String(name);
|
||||
if (key === "class") return this.className || null;
|
||||
if (key === "id") return this.id || null;
|
||||
return this.attributes.has(key) ? this.attributes.get(key) : null;
|
||||
}
|
||||
|
||||
hasAttribute(name) {
|
||||
const key = String(name);
|
||||
if (key === "class") return this._classes.size > 0;
|
||||
if (key === "id") return !!this.id;
|
||||
return this.attributes.has(key);
|
||||
}
|
||||
|
||||
removeAttribute(name) {
|
||||
const key = String(name);
|
||||
this.attributes.delete(key);
|
||||
if (key === "id") this.id = "";
|
||||
if (key === "class") this.className = "";
|
||||
}
|
||||
|
||||
appendChild(child) {
|
||||
if (!child) return child;
|
||||
if (child.parentNode) child.parentNode.removeChild(child);
|
||||
child.parentNode = this;
|
||||
this.children.push(child);
|
||||
return child;
|
||||
}
|
||||
|
||||
insertBefore(child, reference) {
|
||||
if (!child) return child;
|
||||
if (child.parentNode) child.parentNode.removeChild(child);
|
||||
const index = reference ? this.children.indexOf(reference) : -1;
|
||||
child.parentNode = this;
|
||||
if (index < 0) this.children.push(child);
|
||||
else this.children.splice(index, 0, child);
|
||||
return child;
|
||||
}
|
||||
|
||||
removeChild(child) {
|
||||
const index = this.children.indexOf(child);
|
||||
if (index >= 0) this.children.splice(index, 1);
|
||||
child.parentNode = null;
|
||||
return child;
|
||||
}
|
||||
|
||||
remove() {
|
||||
if (this.parentNode) this.parentNode.removeChild(this);
|
||||
}
|
||||
|
||||
contains(node) {
|
||||
let current = node;
|
||||
while (current) {
|
||||
if (current === this) return true;
|
||||
current = current.parentNode;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
matches(selector) {
|
||||
return matchesSelector(this, selector);
|
||||
}
|
||||
|
||||
closest(selector) {
|
||||
let node = this;
|
||||
while (node && node.nodeType === 1) {
|
||||
if (matchesSelector(node, selector)) return node;
|
||||
node = node.parentNode;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
querySelectorAll(selector) {
|
||||
const found = [];
|
||||
const walk = node => {
|
||||
for (const child of node.children) {
|
||||
if (matchesSelector(child, selector)) found.push(child);
|
||||
walk(child);
|
||||
}
|
||||
};
|
||||
walk(this);
|
||||
return found;
|
||||
}
|
||||
|
||||
querySelector(selector) {
|
||||
return this.querySelectorAll(selector)[0] || null;
|
||||
}
|
||||
|
||||
addEventListener(type, callback) {
|
||||
const key = String(type);
|
||||
if (!this._listeners.has(key)) this._listeners.set(key, []);
|
||||
this._listeners.get(key).push(callback);
|
||||
}
|
||||
|
||||
removeEventListener(type, callback) {
|
||||
const key = String(type);
|
||||
const list = this._listeners.get(key);
|
||||
if (!list) return;
|
||||
const index = list.indexOf(callback);
|
||||
if (index >= 0) list.splice(index, 1);
|
||||
}
|
||||
|
||||
dispatchEvent(event) {
|
||||
const type = event?.type;
|
||||
const list = this._listeners.get(type) || [];
|
||||
if (!event.target) event.target = this;
|
||||
for (const callback of [...list]) {
|
||||
callback.call(this, event);
|
||||
}
|
||||
const handler = this[`on${type}`];
|
||||
if (typeof handler === "function") handler.call(this, event);
|
||||
return !event.defaultPrevented;
|
||||
}
|
||||
|
||||
focus() {}
|
||||
select() {}
|
||||
blur() {}
|
||||
click() { this.dispatchEvent({ type: "click", preventDefault() {}, stopPropagation() {} }); }
|
||||
setPointerCapture() {}
|
||||
releasePointerCapture() {}
|
||||
scrollIntoView() {}
|
||||
}
|
||||
|
||||
export function createDocument() {
|
||||
const doc = {
|
||||
nodeType: 9,
|
||||
readyState: "complete",
|
||||
_listeners: new Map(),
|
||||
};
|
||||
doc.body = new FakeElement("body", doc);
|
||||
doc.body._isRoot = true;
|
||||
doc.head = new FakeElement("head", doc);
|
||||
doc.head._isRoot = true;
|
||||
doc.documentElement = new FakeElement("html", doc);
|
||||
doc.documentElement._isRoot = true;
|
||||
doc.documentElement.appendChild(doc.head);
|
||||
doc.documentElement.appendChild(doc.body);
|
||||
|
||||
doc.createElement = tag => new FakeElement(tag, doc);
|
||||
doc.createTextNode = text => ({ nodeType: 3, textContent: String(text), parentNode: null });
|
||||
|
||||
doc.getElementById = id => {
|
||||
const search = node => {
|
||||
if (node.id === id) return node;
|
||||
for (const child of node.children) {
|
||||
const found = search(child);
|
||||
if (found) return found;
|
||||
}
|
||||
return null;
|
||||
};
|
||||
return search(doc.documentElement);
|
||||
};
|
||||
|
||||
doc.querySelectorAll = selector => doc.documentElement.querySelectorAll(selector);
|
||||
doc.querySelector = selector => doc.documentElement.querySelector(selector);
|
||||
|
||||
doc.addEventListener = (type, callback) => {
|
||||
const key = String(type);
|
||||
if (!doc._listeners.has(key)) doc._listeners.set(key, []);
|
||||
doc._listeners.get(key).push(callback);
|
||||
};
|
||||
doc.removeEventListener = (type, callback) => {
|
||||
const list = doc._listeners.get(String(type));
|
||||
if (!list) return;
|
||||
const index = list.indexOf(callback);
|
||||
if (index >= 0) list.splice(index, 1);
|
||||
};
|
||||
doc.dispatchEvent = event => {
|
||||
if (!event.target) event.target = doc.body;
|
||||
for (const callback of [...(doc._listeners.get(event.type) || [])]) {
|
||||
callback.call(doc, event);
|
||||
}
|
||||
return !event.defaultPrevented;
|
||||
};
|
||||
doc.listenerCount = type => (doc._listeners.get(String(type)) || []).length;
|
||||
|
||||
doc.execCommand = () => true;
|
||||
doc.activeElement = doc.body;
|
||||
|
||||
return doc;
|
||||
}
|
||||
|
||||
export function createWindow(doc, overrides = {}) {
|
||||
const win = {
|
||||
innerWidth: 1280,
|
||||
innerHeight: 900,
|
||||
screen: { width: 1920, height: 1080 },
|
||||
isSecureContext: false,
|
||||
location: { origin: "http://localhost", reload() {} },
|
||||
_listeners: new Map(),
|
||||
addEventListener(type, callback) {
|
||||
const key = String(type);
|
||||
if (!win._listeners.has(key)) win._listeners.set(key, []);
|
||||
win._listeners.get(key).push(callback);
|
||||
},
|
||||
removeEventListener(type, callback) {
|
||||
const list = win._listeners.get(String(type));
|
||||
if (!list) return;
|
||||
const index = list.indexOf(callback);
|
||||
if (index >= 0) list.splice(index, 1);
|
||||
},
|
||||
dispatchEvent(event) {
|
||||
for (const callback of [...(win._listeners.get(event.type) || [])]) {
|
||||
callback.call(win, event);
|
||||
}
|
||||
return !event.defaultPrevented;
|
||||
},
|
||||
getComputedStyle(el) {
|
||||
return {
|
||||
display: el?.style?.display || "block",
|
||||
visibility: el?.style?.visibility || "visible",
|
||||
};
|
||||
},
|
||||
requestAnimationFrame(callback) {
|
||||
return setTimeout(() => callback(Date.now()), 0);
|
||||
},
|
||||
cancelAnimationFrame(id) {
|
||||
clearTimeout(id);
|
||||
},
|
||||
listenerCount(type) {
|
||||
return (win._listeners.get(String(type)) || []).length;
|
||||
},
|
||||
...overrides,
|
||||
};
|
||||
return win;
|
||||
}
|
||||
|
||||
export class FakeMutationObserver {
|
||||
constructor(callback) {
|
||||
this.callback = callback;
|
||||
this.target = null;
|
||||
this.options = null;
|
||||
this.disconnected = false;
|
||||
FakeMutationObserver.instances.push(this);
|
||||
}
|
||||
|
||||
observe(target, options) {
|
||||
this.target = target;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
disconnect() {
|
||||
this.disconnected = true;
|
||||
}
|
||||
|
||||
takeRecords() {
|
||||
return [];
|
||||
}
|
||||
|
||||
emit(records) {
|
||||
this.callback(records, this);
|
||||
}
|
||||
}
|
||||
FakeMutationObserver.instances = [];
|
||||
|
||||
export function sleep(ms = 15) {
|
||||
return new Promise(resolve => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
export function stripImports(source) {
|
||||
return source
|
||||
.replace(/^import .*$/gm, "")
|
||||
.replace(/export \{[^}]*\};?\s*$/m, "");
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if test_dir not in sys.path:
|
||||
sys.path.insert(0, test_dir)
|
||||
suite = unittest.defaultTestLoader.discover(test_dir, pattern="test_*.py")
|
||||
runner = unittest.TextTestRunner(verbosity=2)
|
||||
result = runner.run(suite)
|
||||
sys.exit(0 if result.wasSuccessful() else 1)
|
||||
@@ -0,0 +1,136 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.aspect_ratio import TextAspectRatioFinder, PreviewImageSizeAspectRatio
|
||||
|
||||
class TestAspectRatioNodes(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.finder = TextAspectRatioFinder()
|
||||
self.preview_node = PreviewImageSizeAspectRatio()
|
||||
|
||||
def test_find_16_9_standard(self):
|
||||
w, h, ratio = self.finder.find_aspect_ratio(
|
||||
text="A cinematic landscape, 16:9 aspect ratio, 8k",
|
||||
target_mp=1.0,
|
||||
multiple_of=8
|
||||
)
|
||||
self.assertEqual(ratio, "16:9")
|
||||
self.assertGreater(w, h)
|
||||
self.assertEqual(w % 8, 0)
|
||||
self.assertEqual(h % 8, 0)
|
||||
mp = (w * h) / 1_000_000.0
|
||||
self.assertAlmostEqual(mp, 1.0, delta=0.1)
|
||||
|
||||
def test_find_9_16_portrait(self):
|
||||
w, h, ratio = self.finder.find_aspect_ratio(
|
||||
text="Portrait photo, 9:16 vertical format",
|
||||
target_mp=1.0,
|
||||
multiple_of=8
|
||||
)
|
||||
self.assertEqual(ratio, "9:16")
|
||||
self.assertGreater(h, w)
|
||||
self.assertEqual(w % 8, 0)
|
||||
self.assertEqual(h % 8, 0)
|
||||
|
||||
def test_search_mode_front_vs_back(self):
|
||||
text = "Start with 16:9 but override with 1:1 at the end"
|
||||
w_front, h_front, r_front = self.finder.find_aspect_ratio(
|
||||
text=text,
|
||||
search_mode="First match (Front)"
|
||||
)
|
||||
self.assertEqual(r_front, "16:9")
|
||||
|
||||
w_back, h_back, r_back = self.finder.find_aspect_ratio(
|
||||
text=text,
|
||||
search_mode="Last match (Back)"
|
||||
)
|
||||
self.assertEqual(r_back, "1:1")
|
||||
|
||||
def test_fallback_to_default_when_no_ratio_found(self):
|
||||
w, h, ratio = self.finder.find_aspect_ratio(
|
||||
text="No ratio mentioned here at all",
|
||||
default_aspect_ratio="4:3",
|
||||
target_mp=1.0
|
||||
)
|
||||
self.assertEqual(ratio, "4:3")
|
||||
|
||||
def test_edge_case_empty_or_none_text(self):
|
||||
w1, h1, r1 = self.finder.find_aspect_ratio(text="", default_aspect_ratio="1:1")
|
||||
self.assertEqual(r1, "1:1")
|
||||
|
||||
w2, h2, r2 = self.finder.find_aspect_ratio(text=None, default_aspect_ratio="1:1")
|
||||
self.assertEqual(r2, "1:1")
|
||||
|
||||
def test_custom_default_aspect_ratio(self):
|
||||
w, h, ratio = self.finder.find_aspect_ratio(
|
||||
text="No ratio mentioned",
|
||||
default_aspect_ratio="5:4",
|
||||
target_mp=1.0
|
||||
)
|
||||
self.assertEqual(ratio, "5:4")
|
||||
|
||||
w2, h2, ratio2 = self.finder.find_aspect_ratio(
|
||||
text="No ratio mentioned",
|
||||
default_aspect_ratio="30:1",
|
||||
target_mp=1.0
|
||||
)
|
||||
self.assertEqual(ratio2, "30:1")
|
||||
|
||||
def test_corrupt_default_aspect_ratio_falls_back_to_1_1(self):
|
||||
w, h, ratio = self.finder.find_aspect_ratio(
|
||||
text="No ratio mentioned",
|
||||
default_aspect_ratio="corrupted:invalid_ratio",
|
||||
target_mp=1.0
|
||||
)
|
||||
self.assertEqual(ratio, "1:1")
|
||||
|
||||
def test_custom_allowed_ratios_whitelist(self):
|
||||
# When aspect_ratios is configured with "16:9, 9:16", any other ratio in text like "4:3" or "5:7" must be ignored!
|
||||
w, h, ratio = self.finder.find_aspect_ratio(
|
||||
text="First 4:3 then 16:9 and also 5:7",
|
||||
aspect_ratios="16:9, 9:16",
|
||||
default_aspect_ratio="1:1"
|
||||
)
|
||||
self.assertEqual(ratio, "16:9")
|
||||
|
||||
# When prompt only contains unapproved ratios, fall back to default
|
||||
w2, h2, ratio2 = self.finder.find_aspect_ratio(
|
||||
text="A photo in 5:7 format with 12:34 crop",
|
||||
aspect_ratios="16:9, 9:16, 1:1",
|
||||
default_aspect_ratio="1:1"
|
||||
)
|
||||
self.assertEqual(ratio2, "1:1")
|
||||
|
||||
def test_empty_aspect_ratios_accepts_any_ratio(self):
|
||||
# When aspect_ratios is empty, accept any valid aspect ratio found in text
|
||||
w, h, ratio = self.finder.find_aspect_ratio(
|
||||
text="A photo in custom 5:7 format",
|
||||
aspect_ratios="",
|
||||
default_aspect_ratio="1:1"
|
||||
)
|
||||
self.assertEqual(ratio, "5:7")
|
||||
|
||||
w2, h2, ratio2 = self.finder.find_aspect_ratio(
|
||||
text="Anamorphic 2.39:1 movie frame",
|
||||
aspect_ratios=" ",
|
||||
default_aspect_ratio="1:1"
|
||||
)
|
||||
self.assertEqual(ratio2, "2.39:1")
|
||||
|
||||
def test_preview_process(self):
|
||||
with patch("nodes.aspect_ratio.PromptServer.instance.send_sync") as mock_send:
|
||||
res = self.preview_node.process_preview(
|
||||
width=1920,
|
||||
height=1080,
|
||||
unique_id="test_node"
|
||||
)
|
||||
self.assertEqual(res, {})
|
||||
mock_send.assert_called_once()
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,251 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
import copy
|
||||
import shutil
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from PIL import Image, PngImagePlugin
|
||||
import nodes.generation_times as generation_times
|
||||
from nodes.generation_times import (
|
||||
record_generation_time,
|
||||
get_generation_time,
|
||||
patch_prompt_queue,
|
||||
)
|
||||
from nodes.assets_restore import (
|
||||
AssetsRestoreManager,
|
||||
parse_a1111_positive_prompt,
|
||||
_decode_exif_text,
|
||||
)
|
||||
|
||||
|
||||
class TestAssetsRestore(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.output_dir = os.path.join(self.temp_dir, "output")
|
||||
os.makedirs(self.output_dir)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def _make_png(self, name, mtime, prompt_json=None, parameters=None, subfolder=""):
|
||||
folder = os.path.join(self.output_dir, subfolder) if subfolder else self.output_dir
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
path = os.path.join(folder, name)
|
||||
info = PngImagePlugin.PngInfo()
|
||||
if prompt_json is not None:
|
||||
info.add_text("prompt", json.dumps(prompt_json))
|
||||
if parameters is not None:
|
||||
info.add_text("parameters", parameters)
|
||||
Image.new("RGB", (8, 8), (20, 30, 40)).save(path, pnginfo=info)
|
||||
os.utime(path, (mtime, mtime))
|
||||
return path
|
||||
|
||||
def test_scan_returns_newest_limit_and_respects_recursive_setting(self):
|
||||
base = 1_600_000_000
|
||||
self._make_png("old.png", base)
|
||||
self._make_png("new.png", base + 100)
|
||||
self._make_png("nested.png", base + 200, subfolder="child")
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
stats = {}
|
||||
recursive = manager.scan_output_images(self.output_dir, limit=2, recursive=True, stats=stats)
|
||||
self.assertEqual([x["filename"] for x in recursive], ["new.png", "nested.png"])
|
||||
self.assertEqual(stats["scanned"], 3)
|
||||
|
||||
flat = manager.scan_output_images(self.output_dir, limit=10, recursive=False)
|
||||
self.assertEqual([x["filename"] for x in flat], ["old.png", "new.png"])
|
||||
|
||||
def test_extract_png_and_a1111_metadata(self):
|
||||
base = 1_600_000_000
|
||||
prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "hello"}}}
|
||||
self._make_png("meta.png", base, prompt_json=prompt)
|
||||
manager = AssetsRestoreManager()
|
||||
prompt_dict, workflow_dict, positive = manager.extract_image_metadata(
|
||||
os.path.join(self.output_dir, "meta.png"))
|
||||
self.assertEqual(prompt_dict, prompt)
|
||||
self.assertEqual(workflow_dict, {})
|
||||
self.assertEqual(positive, "")
|
||||
|
||||
self._make_png(
|
||||
"a1111.png", base + 1,
|
||||
parameters="A beautiful castle\nNegative prompt: ugly\nSteps: 20, Sampler: Euler")
|
||||
_, _, positive = manager.extract_image_metadata(os.path.join(self.output_dir, "a1111.png"))
|
||||
self.assertEqual(positive, "A beautiful castle")
|
||||
|
||||
def test_decode_exif_user_comment_prefixes(self):
|
||||
# EXIF Unicode comments are UTF-16 big-endian after an 8-byte prefix
|
||||
# (this is what piexif.helper.UserComment.dump produces).
|
||||
unicode_comment = b"UNICODE\x00" + "A castle\nSteps: 20".encode("utf-16-be")
|
||||
self.assertEqual(_decode_exif_text(unicode_comment), "A castle\nSteps: 20")
|
||||
ascii_comment = b"ASCII\x00\x00\x00" + b"plain prompt"
|
||||
self.assertEqual(_decode_exif_text(ascii_comment), "plain prompt")
|
||||
|
||||
def test_parse_a1111_positive_prompt(self):
|
||||
self.assertEqual(
|
||||
parse_a1111_positive_prompt("A cat\n\nNegative prompt: dog\nSteps: 20"),
|
||||
"A cat",
|
||||
)
|
||||
self.assertEqual(
|
||||
parse_a1111_positive_prompt("Just a prompt\nSteps: 10, Sampler: Euler a"),
|
||||
"Just a prompt",
|
||||
)
|
||||
self.assertEqual(parse_a1111_positive_prompt(""), "")
|
||||
|
||||
def test_restore_uses_generation_time_and_builds_debug_report(self):
|
||||
base = 1_600_000_000
|
||||
first = self._make_png("first.png", base, subfolder="sub")
|
||||
second = self._make_png("second.png", base + 60)
|
||||
|
||||
fake_queue = MagicMock()
|
||||
fake_queue.history = {}
|
||||
fake_queue.mutex = threading.RLock()
|
||||
server = MagicMock()
|
||||
server.prompt_queue = fake_queue
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_assets_restore_count", return_value=10), \
|
||||
patch("nodes.assets_restore.get_generation_time", side_effect=lambda fn, sf="", mt=None: 12.5 if fn == "first.png" else None), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
restored = manager.restore_on_launch(server, force=True)
|
||||
|
||||
self.assertEqual(restored, 2)
|
||||
self.assertEqual(len(fake_queue.history), 2)
|
||||
entries = list(fake_queue.history.values())
|
||||
self.assertEqual(entries[0]["outputs"]["0"]["images"][0]["filename"], "first.png")
|
||||
start_ts = entries[0]["status"]["messages"][0][1]["timestamp"]
|
||||
end_ts = entries[0]["status"]["messages"][1][1]["timestamp"]
|
||||
self.assertEqual(end_ts - start_ts, 12500)
|
||||
self.assertEqual(entries[0]["status"]["saturnnodes_duration"], 12.5)
|
||||
self.assertEqual(entries[0]["prompt"][3]["extra_pnginfo"]["saturnnodes_generation_time"], 12.5)
|
||||
self.assertEqual(len(entries[0]["status"]["messages"]), 2)
|
||||
# No recorded duration: keep the mtime for sorting, but do not emit an
|
||||
# execution_success timestamp that would render as a fake "0.00s".
|
||||
self.assertEqual(entries[1]["status"]["messages"][0][1]["timestamp"], base * 1000 + 60 * 1000)
|
||||
self.assertEqual(len(entries[1]["status"]["messages"]), 1)
|
||||
|
||||
report = manager.last_debug_report
|
||||
self.assertEqual(report["restored_count"], 2)
|
||||
self.assertEqual(report["generation_time_matches"], 1)
|
||||
self.assertEqual(report["total_scanned_images"], 2)
|
||||
self.assertIn("elapsed_ms", report)
|
||||
self.assertTrue(os.path.exists(first) and os.path.exists(second))
|
||||
|
||||
|
||||
class TestGenerationTimes(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.old_file = generation_times.GENERATION_TIMES_FILE
|
||||
self.old_data = generation_times._DATA
|
||||
self.old_mtime = generation_times._DATA_MTIME
|
||||
self.old_patched = generation_times._PATCHED
|
||||
self.old_starts = dict(generation_times._START_TIMES)
|
||||
generation_times.GENERATION_TIMES_FILE = os.path.join(self.temp_dir, "generation_times.json")
|
||||
generation_times._DATA = None
|
||||
generation_times._DATA_MTIME = None
|
||||
generation_times._PATCHED = False
|
||||
generation_times._START_TIMES.clear()
|
||||
|
||||
def tearDown(self):
|
||||
generation_times.GENERATION_TIMES_FILE = self.old_file
|
||||
generation_times._DATA = self.old_data
|
||||
generation_times._DATA_MTIME = self.old_mtime
|
||||
generation_times._PATCHED = self.old_patched
|
||||
generation_times._START_TIMES.clear()
|
||||
generation_times._START_TIMES.update(self.old_starts)
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_record_and_lookup_generation_time(self):
|
||||
record_generation_time("pid-1", 3.25, [
|
||||
{"filename": "a.png", "subfolder": ""},
|
||||
{"filename": "b.png", "subfolder": "nested"},
|
||||
])
|
||||
self.assertEqual(get_generation_time("a.png"), 3.25)
|
||||
self.assertEqual(get_generation_time("b.png", "nested"), 3.25)
|
||||
self.assertEqual(get_generation_time("b.png"), 3.25) # bare filename fallback
|
||||
self.assertIsNone(get_generation_time("missing.png"))
|
||||
|
||||
with open(generation_times.GENERATION_TIMES_FILE, "r", encoding="utf-8") as f:
|
||||
raw = f.read()
|
||||
self.assertNotIn("\n", raw) # compact JSON
|
||||
|
||||
def test_moved_file_fallback_is_timestamp_guarded_and_unambiguous(self):
|
||||
now = time.time()
|
||||
record_generation_time("pid-moved", 7.0, [
|
||||
{"filename": "moved.png", "subfolder": "old"},
|
||||
])
|
||||
# A moved file has a new subfolder/key; the bare-name fallback works
|
||||
# only when the recorded completion time is close to the file mtime.
|
||||
self.assertEqual(get_generation_time("moved.png", "new", now), 7.0)
|
||||
self.assertIsNone(get_generation_time("moved.png", "new", now - 3600))
|
||||
|
||||
record_generation_time("pid-other", 9.0, [
|
||||
{"filename": "moved.png", "subfolder": "elsewhere"},
|
||||
])
|
||||
# Two candidate records are ambiguous, so no duration is invented.
|
||||
self.assertIsNone(get_generation_time("moved.png", "new", now))
|
||||
|
||||
def test_temporary_outputs_are_not_recorded_for_assets(self):
|
||||
files = generation_times._extract_output_files({
|
||||
"outputs": {"1": {"images": [
|
||||
{"filename": "preview.png", "subfolder": "", "type": "temp"},
|
||||
{"filename": "saved.png", "subfolder": "", "type": "output"},
|
||||
{"filename": "no-type.png", "subfolder": ""},
|
||||
]}}
|
||||
})
|
||||
self.assertEqual([f["filename"] for f in files], ["saved.png", "no-type.png"])
|
||||
record_generation_time("pid-temp", 2.0, files)
|
||||
self.assertIsNone(get_generation_time("preview.png"))
|
||||
self.assertEqual(get_generation_time("saved.png"), 2.0)
|
||||
|
||||
def test_queue_patch_measures_task_duration(self):
|
||||
class FakeQueue:
|
||||
def __init__(self):
|
||||
self.currently_running = {}
|
||||
self.mutex = threading.RLock()
|
||||
self.items = [(0, "pid-42", {"prompt": {}}, {}, {}, True)]
|
||||
self.task_counter = 0
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
if not self.items:
|
||||
return None
|
||||
item = self.items.pop(0)
|
||||
task_id = self.task_counter
|
||||
self.task_counter += 1
|
||||
self.currently_running[task_id] = copy.deepcopy(item)
|
||||
# Match ComfyUI's real PromptQueue.get return shape.
|
||||
return (item, task_id)
|
||||
|
||||
def task_done(self, item_id, history_result, status, process_item=None):
|
||||
self.currently_running.pop(item_id, None)
|
||||
|
||||
queue = FakeQueue()
|
||||
server = MagicMock()
|
||||
server.prompt_queue = queue
|
||||
|
||||
patch_prompt_queue(server)
|
||||
item, task_id = queue.get()
|
||||
queue.task_done(task_id, {
|
||||
"outputs": {"9": {"images": [{"filename": "queued.png", "subfolder": "run"}]}}
|
||||
}, "success")
|
||||
|
||||
duration = get_generation_time("queued.png", "run")
|
||||
self.assertIsNotNone(duration)
|
||||
self.assertGreaterEqual(duration, 0.0)
|
||||
with open(generation_times.GENERATION_TIMES_FILE, "r", encoding="utf-8") as f:
|
||||
saved = json.load(f)
|
||||
self.assertEqual(saved["entries"]["run/queued.png"]["status"], "success")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,422 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper # noqa: F401 (redirects SaturnNodes state to a temp dir)
|
||||
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
from nodes.assets_restore import (
|
||||
AssetsRestoreManager,
|
||||
_decode_exif_text,
|
||||
get_assets_restore_count,
|
||||
parse_a1111_positive_prompt,
|
||||
)
|
||||
|
||||
|
||||
class TestAssetsRestoreExtended(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.output_dir = os.path.join(self.temp_dir, "output")
|
||||
os.makedirs(self.output_dir)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def _make_png(self, name, mtime, subfolder="", prompt_json=None, parameters=None):
|
||||
folder = os.path.join(self.output_dir, subfolder) if subfolder else self.output_dir
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
path = os.path.join(folder, name)
|
||||
info = PngImagePlugin.PngInfo()
|
||||
if prompt_json is not None:
|
||||
info.add_text("prompt", json.dumps(prompt_json))
|
||||
if parameters is not None:
|
||||
info.add_text("parameters", parameters)
|
||||
Image.new("RGB", (8, 8), (10, 20, 30)).save(path, pnginfo=info)
|
||||
os.utime(path, (mtime, mtime))
|
||||
return path
|
||||
|
||||
def _fake_queue(self, history=None, mutex=True):
|
||||
queue = SimpleNamespace(history=history)
|
||||
if mutex:
|
||||
queue.mutex = threading.RLock()
|
||||
return queue
|
||||
|
||||
def _server(self, queue):
|
||||
return SimpleNamespace(prompt_queue=queue, queue_updated=lambda: None)
|
||||
|
||||
def test_scan_equal_mtimes_use_path_tiebreak(self):
|
||||
base = 1_700_000_000
|
||||
# Deliberately create in non-sorted order so filesystem iteration order
|
||||
# cannot accidentally make the test pass.
|
||||
self._make_png("b.png", base)
|
||||
self._make_png("a.png", base)
|
||||
self._make_png("c.png", base, subfolder="sub")
|
||||
self._make_png("a.png", base, subfolder="sub")
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
first = manager.scan_output_images(self.output_dir, limit=10)
|
||||
second = manager.scan_output_images(self.output_dir, limit=10)
|
||||
|
||||
expected = [
|
||||
("a.png", ""),
|
||||
("b.png", ""),
|
||||
("a.png", "sub"),
|
||||
("c.png", "sub"),
|
||||
]
|
||||
self.assertEqual([(x["filename"], x["subfolder"]) for x in first], expected)
|
||||
self.assertEqual(first, second, "repeated scans must return identical ordering")
|
||||
|
||||
def test_scan_limit_is_clamped_and_invalid_limits_fall_back(self):
|
||||
base = 1_700_000_000
|
||||
for index in range(4):
|
||||
self._make_png(f"img{index}.png", base + index)
|
||||
with open(os.path.join(self.output_dir, "notes.txt"), "w", encoding="utf-8") as handle:
|
||||
handle.write("not an image")
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
stats = {}
|
||||
newest_one = manager.scan_output_images(self.output_dir, limit=0, stats=stats)
|
||||
self.assertEqual([x["filename"] for x in newest_one], ["img3.png"])
|
||||
self.assertEqual(stats["scanned"], 4, "non-image files must not count as scanned images")
|
||||
|
||||
self.assertEqual(
|
||||
[x["filename"] for x in manager.scan_output_images(self.output_dir, limit=-7)],
|
||||
["img3.png"],
|
||||
)
|
||||
self.assertEqual(
|
||||
len(manager.scan_output_images(self.output_dir, limit="not-a-number")),
|
||||
4,
|
||||
"an invalid limit must fall back to the default scan size",
|
||||
)
|
||||
|
||||
def test_restore_limit_is_clamped_and_invalid_uses_setting(self):
|
||||
base = 1_700_000_000
|
||||
self._make_png("one.png", base)
|
||||
self._make_png("two.png", base + 1)
|
||||
self._make_png("three.png", base + 2)
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
self.assertEqual(manager.restore_on_launch(self._server(queue), limit=0, force=True), 1)
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_assets_restore_count", return_value=2), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
self.assertEqual(
|
||||
manager.restore_on_launch(self._server(queue), limit="broken", force=True),
|
||||
2,
|
||||
)
|
||||
|
||||
def test_concurrent_force_restores_are_serialized(self):
|
||||
self._make_png("image.png", 1_700_000_000)
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
server = self._server(queue)
|
||||
|
||||
intervals = []
|
||||
intervals_lock = threading.Lock()
|
||||
|
||||
def slow_ingest(prepared):
|
||||
start = time.monotonic()
|
||||
time.sleep(0.08)
|
||||
end = time.monotonic()
|
||||
with intervals_lock:
|
||||
intervals.append((start, end))
|
||||
return len(prepared)
|
||||
|
||||
errors = []
|
||||
|
||||
def run_restore():
|
||||
try:
|
||||
manager.restore_on_launch(server, force=True)
|
||||
except Exception as exc: # pragma: no cover - surfaced via assert below
|
||||
errors.append(exc)
|
||||
|
||||
# Apply the module patches in the main thread. Patching the same target
|
||||
# from two threads can restore stale values when the contexts unwind in
|
||||
# a different order and would make this test flaky.
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", side_effect=slow_ingest):
|
||||
threads = [threading.Thread(target=run_restore) for _ in range(2)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join(10)
|
||||
|
||||
self.assertEqual(errors, [])
|
||||
self.assertEqual(len(intervals), 2)
|
||||
intervals.sort()
|
||||
self.assertGreaterEqual(
|
||||
intervals[1][0],
|
||||
intervals[0][1] - 0.01,
|
||||
"the restore lock must prevent overlapping restore/ingest runs",
|
||||
)
|
||||
|
||||
def test_restore_preserves_cross_day_timestamps_and_durations(self):
|
||||
day = 86_400
|
||||
base = 1_700_000_000
|
||||
specs = [
|
||||
("day-one.png", base, 1.5),
|
||||
("day-two.png", base + day, 2.0),
|
||||
("day-three.png", base + 2 * day, 3.25),
|
||||
]
|
||||
durations = {}
|
||||
for name, mtime, duration in specs:
|
||||
self._make_png(name, mtime)
|
||||
durations[name] = duration
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch(
|
||||
"nodes.assets_restore.get_generation_time",
|
||||
side_effect=lambda fn, sf="", mt=None: durations.get(fn),
|
||||
), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
restored = manager.restore_on_launch(server, force=True)
|
||||
|
||||
self.assertEqual(restored, 3)
|
||||
entries = list(queue.history.values())
|
||||
self.assertEqual(
|
||||
[entry["prompt"][0] for entry in entries],
|
||||
[1, 2, 3],
|
||||
"restored history must be numbered oldest to newest across days",
|
||||
)
|
||||
|
||||
for entry, (name, mtime, duration) in zip(entries, specs):
|
||||
mtime_ms = int(mtime * 1000)
|
||||
duration_ms = int(duration * 1000)
|
||||
self.assertEqual(entry["prompt"][3]["create_time"], mtime_ms, name)
|
||||
messages = entry["status"]["messages"]
|
||||
self.assertEqual([message[0] for message in messages], ["execution_start", "execution_success"])
|
||||
self.assertEqual(messages[0][1]["timestamp"], mtime_ms - duration_ms, name)
|
||||
self.assertEqual(messages[1][1]["timestamp"], mtime_ms, name)
|
||||
self.assertEqual(entry["status"]["saturnnodes_duration"], duration, name)
|
||||
self.assertEqual(
|
||||
entry["prompt"][3]["extra_pnginfo"]["saturnnodes_generation_time"],
|
||||
duration,
|
||||
name,
|
||||
)
|
||||
self.assertEqual(
|
||||
entry["outputs"]["0"]["images"][0]["filename"],
|
||||
name,
|
||||
)
|
||||
|
||||
def test_restore_recovers_from_corrupt_image(self):
|
||||
path = os.path.join(self.output_dir, "broken.png")
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(b"this is not a PNG")
|
||||
os.utime(path, (1_700_000_000, 1_700_000_000))
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
restored = manager.restore_on_launch(server, force=True)
|
||||
|
||||
self.assertEqual(restored, 1)
|
||||
entry = next(iter(queue.history.values()))
|
||||
self.assertEqual(entry["prompt"][2], {}, "corrupt metadata must fall back to an empty prompt")
|
||||
self.assertEqual(entry["outputs"]["0"]["images"][0]["filename"], "broken.png")
|
||||
self.assertEqual(entry["status"]["status_str"], "success")
|
||||
self.assertEqual(entry["status"]["messages"], [["execution_start", {"timestamp": 1_700_000_000_000}]])
|
||||
|
||||
def test_restore_handles_missing_mutex_and_bad_history(self):
|
||||
self._make_png("image.png", 1_700_000_000)
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue(history=None, mutex=False)
|
||||
server = self._server(queue)
|
||||
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
restored = manager.restore_on_launch(server, force=True)
|
||||
|
||||
self.assertEqual(restored, 1)
|
||||
self.assertIsInstance(queue.history, dict)
|
||||
self.assertEqual(len(queue.history), 1)
|
||||
|
||||
def test_restore_survives_asset_db_ingest_failure(self):
|
||||
self._make_png("image.png", 1_700_000_000)
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", side_effect=RuntimeError("asset db unavailable")):
|
||||
restored = manager.restore_on_launch(server, force=True)
|
||||
|
||||
self.assertEqual(restored, 1)
|
||||
self.assertEqual(manager.last_debug_report["asset_db_registered"], 0)
|
||||
self.assertEqual(manager.last_debug_report["status"], "success")
|
||||
|
||||
def test_restore_ignores_non_finite_or_negative_durations(self):
|
||||
base = 1_700_000_000
|
||||
for name in ("nan.png", "inf.png", "neg.png", "ok.png"):
|
||||
self._make_png(name, base)
|
||||
base += 1
|
||||
|
||||
durations = {
|
||||
"nan.png": float("nan"),
|
||||
"inf.png": float("inf"),
|
||||
"neg.png": -2.5,
|
||||
"ok.png": 4.0,
|
||||
}
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", side_effect=lambda fn, sf="", mt=None: durations[fn]), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
restored = manager.restore_on_launch(server, force=True)
|
||||
|
||||
self.assertEqual(restored, 4)
|
||||
entries = {
|
||||
entry["outputs"]["0"]["images"][0]["filename"]: entry
|
||||
for entry in queue.history.values()
|
||||
}
|
||||
for name in ("nan.png", "inf.png", "neg.png"):
|
||||
self.assertEqual(len(entries[name]["status"]["messages"]), 1, name)
|
||||
self.assertNotIn("saturnnodes_duration", entries[name]["status"], name)
|
||||
self.assertNotIn(
|
||||
"saturnnodes_generation_time",
|
||||
entries[name]["prompt"][3]["extra_pnginfo"],
|
||||
name,
|
||||
)
|
||||
self.assertEqual(len(entries["ok.png"]["status"]["messages"]), 2)
|
||||
self.assertEqual(entries["ok.png"]["status"]["saturnnodes_duration"], 4.0)
|
||||
|
||||
def test_restore_disabled_marks_launch_and_skips_scan(self):
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch("nodes.assets_restore.is_assets_restore_enabled", return_value=False), \
|
||||
patch.object(manager, "scan_output_images") as scan:
|
||||
self.assertEqual(manager.restore_on_launch(server), 0)
|
||||
self.assertEqual(manager.restore_on_launch(server), 0)
|
||||
|
||||
scan.assert_not_called()
|
||||
self.assertTrue(manager._restored_on_launch)
|
||||
|
||||
def test_restore_runs_once_without_force(self):
|
||||
self._make_png("image.png", 1_700_000_000)
|
||||
manager = AssetsRestoreManager()
|
||||
queue = self._fake_queue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0), \
|
||||
patch.object(manager, "scan_output_images", wraps=manager.scan_output_images) as scan:
|
||||
self.assertEqual(manager.restore_on_launch(server), 1)
|
||||
self.assertEqual(manager.restore_on_launch(server), 0)
|
||||
|
||||
self.assertEqual(scan.call_count, 1)
|
||||
|
||||
def test_restored_prompt_ids_are_stable_and_unique(self):
|
||||
base = 1_700_000_000
|
||||
self._make_png("same.png", base)
|
||||
self._make_png("same.png", base + 1, subfolder="one")
|
||||
self._make_png("same.png", base + 2, subfolder="two")
|
||||
|
||||
manager = AssetsRestoreManager()
|
||||
first_queue = self._fake_queue()
|
||||
second_queue = self._fake_queue()
|
||||
|
||||
for queue in (first_queue, second_queue):
|
||||
with patch("nodes.assets_restore.folder_paths.get_output_directory", return_value=self.output_dir), \
|
||||
patch("nodes.assets_restore.is_assets_restore_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.is_recursive_assets_scan_enabled", return_value=True), \
|
||||
patch("nodes.assets_restore.get_generation_time", return_value=None), \
|
||||
patch.object(manager, "_ingest_assets", return_value=0):
|
||||
manager.restore_on_launch(self._server(queue), force=True)
|
||||
|
||||
first_ids = set(first_queue.history)
|
||||
second_ids = set(second_queue.history)
|
||||
self.assertEqual(first_ids, second_ids)
|
||||
self.assertEqual(len(first_ids), 3)
|
||||
self.assertEqual(
|
||||
first_ids,
|
||||
{
|
||||
str(uuid.uuid5(uuid.NAMESPACE_URL, f"saturnnodes:{key}"))
|
||||
for key in ("same.png", "one/same.png", "two/same.png")
|
||||
},
|
||||
)
|
||||
for entry in first_queue.history.values():
|
||||
self.assertIn(entry["prompt"][1], first_ids)
|
||||
|
||||
def test_decode_exif_text_strips_bom_and_nul_padding(self):
|
||||
unicode_comment = b"UNICODE\x00" + "\ufeffA castle\x00".encode("utf-16-be")
|
||||
self.assertEqual(_decode_exif_text(unicode_comment), "A castle")
|
||||
self.assertEqual(
|
||||
_decode_exif_text(b"ASCII\x00\x00\x00" + b"A\x00B\x00"),
|
||||
"AB",
|
||||
)
|
||||
self.assertEqual(_decode_exif_text("\ufeff padded \x00"), "padded")
|
||||
self.assertEqual(_decode_exif_text(None), "")
|
||||
self.assertEqual(_decode_exif_text(b""), "")
|
||||
|
||||
def test_parse_a1111_accepts_bytes_and_crlf(self):
|
||||
raw = b"A cat\r\nNegative prompt: dog\r\nSteps: 20, Sampler: Euler"
|
||||
self.assertEqual(parse_a1111_positive_prompt(raw), "A cat")
|
||||
self.assertEqual(
|
||||
parse_a1111_positive_prompt(b"Line one\r\nLine two\r\nSteps: 4"),
|
||||
"Line one\nLine two",
|
||||
)
|
||||
self.assertEqual(parse_a1111_positive_prompt(None), "")
|
||||
|
||||
def test_generation_time_count_setting_clamps(self):
|
||||
with patch("nodes.assets_restore.get_env_setting", return_value="0"):
|
||||
self.assertEqual(get_assets_restore_count(), 1)
|
||||
with patch("nodes.assets_restore.get_env_setting", return_value="100000"):
|
||||
self.assertEqual(get_assets_restore_count(), 1000)
|
||||
with patch("nodes.assets_restore.get_env_setting", return_value="not-a-number"):
|
||||
self.assertEqual(get_assets_restore_count(), 64)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,97 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import tempfile
|
||||
import shutil
|
||||
from PIL import Image
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.auto_watcher import LoadImageFromFolder
|
||||
|
||||
class TestAutoWatcher(unittest.TestCase):
|
||||
def setUp(self):
|
||||
import folder_paths
|
||||
self.node = LoadImageFromFolder()
|
||||
inp_dir = folder_paths.get_input_directory()
|
||||
os.makedirs(inp_dir, exist_ok=True)
|
||||
self.temp_dir = tempfile.mkdtemp(dir=inp_dir)
|
||||
self.rel_dir = os.path.relpath(self.temp_dir, inp_dir)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_create_dummy_image(self):
|
||||
dummy = self.node.create_dummy_image()
|
||||
self.assertEqual(dummy.shape, (1, 512, 512, 3))
|
||||
|
||||
def test_get_filtered_files_by_regex(self):
|
||||
img = Image.new("RGB", (64, 64), color="blue")
|
||||
f1 = os.path.join(self.temp_dir, "test_01.png")
|
||||
f2 = os.path.join(self.temp_dir, "test_02.jpg")
|
||||
f3 = os.path.join(self.temp_dir, "other.png")
|
||||
f4 = os.path.join(self.temp_dir, "ignored.txt")
|
||||
img.save(f1)
|
||||
img.save(f2)
|
||||
img.save(f3)
|
||||
with open(f4, "w") as f:
|
||||
f.write("text file")
|
||||
|
||||
files = self.node.get_filtered_files(self.temp_dir, sort_by="name", regex_filter=r"test_\d+")
|
||||
self.assertEqual(len(files), 2)
|
||||
self.assertIn(f1, files)
|
||||
self.assertIn(f2, files)
|
||||
self.assertNotIn(f3, files)
|
||||
self.assertNotIn(f4, files)
|
||||
|
||||
def test_empty_folder_returns_dummy(self):
|
||||
empty_dir = os.path.join(self.temp_dir, "empty")
|
||||
os.makedirs(empty_dir, exist_ok=True)
|
||||
empty_rel = os.path.join(self.rel_dir, "empty")
|
||||
img_tensor, has_image = self.node.watch(
|
||||
folder=empty_rel,
|
||||
wait_if_folder_is_empty=False,
|
||||
rescan_interval=1,
|
||||
sort_by="date_modified",
|
||||
regex_filter=".*"
|
||||
)
|
||||
self.assertFalse(has_image)
|
||||
self.assertEqual(img_tensor.shape, (1, 512, 512, 3))
|
||||
|
||||
def test_sequential_cycle_when_delete_false(self):
|
||||
"""Verify that when delete_image=False, consecutive runs cycle through files instead of loading the same file forever."""
|
||||
f1 = os.path.join(self.temp_dir, "img_a.png")
|
||||
f2 = os.path.join(self.temp_dir, "img_b.png")
|
||||
Image.new("RGB", (64, 64), color="red").save(f1)
|
||||
Image.new("RGB", (64, 64), color="green").save(f2)
|
||||
|
||||
# Run 1 -> loads img_a
|
||||
t1, has1 = self.node.watch(folder=self.rel_dir, wait_if_folder_is_empty=False, rescan_interval=1, sort_by="name", regex_filter=".*", delete_image=False, unique_id="cycle_test")
|
||||
self.assertTrue(has1)
|
||||
|
||||
# Run 2 -> loads img_b (cycled!)
|
||||
t2, has2 = self.node.watch(folder=self.rel_dir, wait_if_folder_is_empty=False, rescan_interval=1, sort_by="name", regex_filter=".*", delete_image=False, unique_id="cycle_test")
|
||||
self.assertTrue(has2)
|
||||
|
||||
# Run 3 -> loops back to img_a
|
||||
t3, has3 = self.node.watch(folder=self.rel_dir, wait_if_folder_is_empty=False, rescan_interval=1, sort_by="name", regex_filter=".*", delete_image=False, unique_id="cycle_test")
|
||||
self.assertTrue(has3)
|
||||
|
||||
def test_load_and_remove_image_with_delete_toggle(self):
|
||||
img = Image.new("RGB", (64, 64), color="red")
|
||||
img_path = os.path.join(self.temp_dir, "test_del.png")
|
||||
img.save(img_path)
|
||||
|
||||
# 1. delete_image = False (default: file must remain on disk)
|
||||
tensor = self.node.load_and_remove_image(img_path, delete_image=False)
|
||||
self.assertEqual(tensor.shape, (1, 64, 64, 3))
|
||||
self.assertTrue(os.path.exists(img_path), "File should not be deleted when delete_image=False")
|
||||
|
||||
# 2. delete_image = True (file must be removed from disk)
|
||||
tensor2 = self.node.load_and_remove_image(img_path, delete_image=True)
|
||||
self.assertEqual(tensor2.shape, (1, 64, 64, 3))
|
||||
self.assertFalse(os.path.exists(img_path), "File should be deleted when delete_image=True")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,337 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import json
|
||||
import tempfile
|
||||
import threading
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.queue_control import PersistentQueueManager, PauseQueueManager
|
||||
import nodes.queue_control as queue_control
|
||||
|
||||
# 36 Colors matching web/batch_queue.js
|
||||
BATCH_COLORS = [
|
||||
"#10b981", "#3b82f6", "#8b5cf6", "#ec4899", "#f59e0b", "#06b6d4",
|
||||
"#ef4444", "#14b8a6", "#6366f1", "#f97316", "#84cc16", "#d946ef",
|
||||
"#0ea5e9", "#a855f7", "#e11d48", "#22c55e", "#eab308", "#64748b",
|
||||
"#2dd4bf", "#fb7185", "#38bdf8", "#c084fc", "#f43f5e", "#4ade80",
|
||||
"#fbbf24", "#818cf8", "#fb923c", "#a3e635", "#f472b6", "#22d3ee",
|
||||
"#e879f9", "#a7f3d0", "#fde047", "#fda4af", "#93c5fd", "#c4b5fd"
|
||||
]
|
||||
|
||||
def calculate_segment_types(ordered_pids, registry):
|
||||
"""
|
||||
Python mirror of the segment calculation logic in web/batch_queue.js
|
||||
"""
|
||||
segments = {}
|
||||
for pid in ordered_pids:
|
||||
info = registry.get(pid)
|
||||
if not info or not info.get("batchId"):
|
||||
segments[pid] = None
|
||||
continue
|
||||
|
||||
batch_id = info["batchId"]
|
||||
occurrences = [p for p in ordered_pids if registry.get(p, {}).get("batchId") == batch_id]
|
||||
|
||||
if len(occurrences) <= 1:
|
||||
batch_count = info.get("batchCount", 1)
|
||||
if batch_count == 1:
|
||||
segments[pid] = "single"
|
||||
else:
|
||||
item_index = info.get("itemIndex", 0)
|
||||
if item_index == 0:
|
||||
segments[pid] = "first"
|
||||
elif item_index == (batch_count - 1):
|
||||
segments[pid] = "last"
|
||||
else:
|
||||
segments[pid] = "middle"
|
||||
continue
|
||||
|
||||
first_pid = occurrences[0]
|
||||
last_pid = occurrences[-1]
|
||||
|
||||
if pid == first_pid:
|
||||
segments[pid] = "first"
|
||||
elif pid == last_pid:
|
||||
segments[pid] = "last"
|
||||
else:
|
||||
segments[pid] = "middle"
|
||||
|
||||
return segments
|
||||
|
||||
|
||||
class TestBatchQueueLogic(unittest.TestCase):
|
||||
def test_36_color_cycling(self):
|
||||
self.assertEqual(len(BATCH_COLORS), 36)
|
||||
|
||||
# Batch 1 -> color index 0
|
||||
self.assertEqual((1 - 1) % 36, 0)
|
||||
self.assertEqual(BATCH_COLORS[(1 - 1) % 36], "#10b981")
|
||||
|
||||
# Batch 36 -> color index 35
|
||||
self.assertEqual((36 - 1) % 36, 35)
|
||||
self.assertEqual(BATCH_COLORS[(36 - 1) % 36], "#c4b5fd")
|
||||
|
||||
# Batch 37 -> cycles back to color index 0
|
||||
self.assertEqual((37 - 1) % 36, 0)
|
||||
self.assertEqual(BATCH_COLORS[(37 - 1) % 36], "#10b981")
|
||||
|
||||
# Batch 73 -> cycles back to color index 0
|
||||
self.assertEqual((73 - 1) % 36, 0)
|
||||
|
||||
def test_single_item_batch(self):
|
||||
registry = {
|
||||
"p1": {"batchId": "b1", "batchCount": 1, "itemIndex": 0}
|
||||
}
|
||||
segments = calculate_segment_types(["p1"], registry)
|
||||
self.assertEqual(segments["p1"], "single")
|
||||
|
||||
def test_intact_contiguous_batch(self):
|
||||
registry = {
|
||||
"p1": {"batchId": "b1", "batchCount": 3, "itemIndex": 0},
|
||||
"p2": {"batchId": "b1", "batchCount": 3, "itemIndex": 1},
|
||||
"p3": {"batchId": "b1", "batchCount": 3, "itemIndex": 2},
|
||||
}
|
||||
segments = calculate_segment_types(["p1", "p2", "p3"], registry)
|
||||
self.assertEqual(segments["p1"], "first")
|
||||
self.assertEqual(segments["p2"], "middle")
|
||||
self.assertEqual(segments["p3"], "last")
|
||||
|
||||
def test_interrupted_batch_inserted_in_between(self):
|
||||
# Batch A was [pA1, pA2], but prompt B was inserted in-between: [pA1, pB, pA2]
|
||||
registry = {
|
||||
"pA1": {"batchId": "bA", "batchCount": 2, "itemIndex": 0},
|
||||
"pB": {"batchId": "bB", "batchCount": 1, "itemIndex": 0},
|
||||
"pA2": {"batchId": "bA", "batchCount": 2, "itemIndex": 1},
|
||||
}
|
||||
segments = calculate_segment_types(["pA1", "pB", "pA2"], registry)
|
||||
# pA1 is the start of Batch A (curves top, open bottom)
|
||||
self.assertEqual(segments["pA1"], "first")
|
||||
# pB is a standalone single item (curves top and bottom)
|
||||
self.assertEqual(segments["pB"], "single")
|
||||
# pA2 is the end of Batch A (open top, curves bottom)
|
||||
self.assertEqual(segments["pA2"], "last")
|
||||
|
||||
def test_multiple_interruptions(self):
|
||||
# [pA1, pB, pA2, pC, pA3]
|
||||
registry = {
|
||||
"pA1": {"batchId": "bA", "batchCount": 3, "itemIndex": 0},
|
||||
"pB": {"batchId": "bB", "batchCount": 1, "itemIndex": 0},
|
||||
"pA2": {"batchId": "bA", "batchCount": 3, "itemIndex": 1},
|
||||
"pC": {"batchId": "bC", "batchCount": 1, "itemIndex": 0},
|
||||
"pA3": {"batchId": "bA", "batchCount": 3, "itemIndex": 2},
|
||||
}
|
||||
segments = calculate_segment_types(["pA1", "pB", "pA2", "pC", "pA3"], registry)
|
||||
self.assertEqual(segments["pA1"], "first")
|
||||
self.assertEqual(segments["pB"], "single")
|
||||
self.assertEqual(segments["pA2"], "middle")
|
||||
self.assertEqual(segments["pC"], "single")
|
||||
self.assertEqual(segments["pA3"], "last")
|
||||
|
||||
|
||||
class TestPersistentQueueBatchMetadata(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.TemporaryDirectory()
|
||||
self.temp_file = os.path.join(self.temp_dir.name, "test_persistent_queue.json")
|
||||
self.patcher = patch("nodes.queue_control.PERSISTENT_FILE", self.temp_file)
|
||||
self.patcher.start()
|
||||
self.enabled_patcher = patch("nodes.queue_control.is_persistent_queue_enabled", return_value=True)
|
||||
self.enabled_patcher.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.enabled_patcher.stop()
|
||||
self.patcher.stop()
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
def test_save_and_restore_batch_info(self):
|
||||
mgr = PersistentQueueManager()
|
||||
|
||||
# Set batch info for a prompt
|
||||
batch_data = {
|
||||
"batchId": "b_12345",
|
||||
"batchNumber": 5,
|
||||
"batchCount": 3,
|
||||
"itemIndex": 0,
|
||||
"color": "#10b981",
|
||||
"colorIndex": 4
|
||||
}
|
||||
mgr.set_batch_info("prompt_99", batch_data)
|
||||
self.assertIn("prompt_99", mgr.batch_meta)
|
||||
|
||||
# Add item to persistent queue
|
||||
item_tuple = (0, "prompt_99", {"prompt": {}}, {}, {})
|
||||
mgr.add_item(item_tuple)
|
||||
mgr.flush()
|
||||
|
||||
# Verify batch_info is attached inside persistent_items
|
||||
saved_entry = [x for x in mgr.persistent_items if x["prompt_id"] == "prompt_99"][0]
|
||||
self.assertEqual(saved_entry.get("batch_info"), batch_data)
|
||||
|
||||
# Simulate reload from disk
|
||||
new_mgr = PersistentQueueManager()
|
||||
self.assertIn("prompt_99", new_mgr.batch_meta)
|
||||
self.assertEqual(new_mgr.batch_meta["prompt_99"]["batchId"], "b_12345")
|
||||
|
||||
def test_remove_item_cleans_batch_meta(self):
|
||||
mgr = PersistentQueueManager()
|
||||
mgr.set_batch_info("p1", {"batchId": "b1"})
|
||||
mgr.add_item((0, "p1", {}, {}, {}))
|
||||
self.assertIn("p1", mgr.batch_meta)
|
||||
|
||||
mgr.remove_item("p1")
|
||||
mgr.flush()
|
||||
self.assertNotIn("p1", mgr.batch_meta)
|
||||
self.assertEqual(len(mgr.persistent_items), 0)
|
||||
|
||||
def test_wipe_all_cleans_batch_meta(self):
|
||||
mgr = PersistentQueueManager()
|
||||
mgr.set_batch_info("p1", {"batchId": "b1"})
|
||||
mgr.set_batch_info("p2", {"batchId": "b2"})
|
||||
mgr.wipe_all()
|
||||
self.assertEqual(len(mgr.batch_meta), 0)
|
||||
self.assertEqual(len(mgr.persistent_items), 0)
|
||||
|
||||
|
||||
class TestPersistentQueueHardening(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.TemporaryDirectory()
|
||||
self.temp_file = os.path.join(self.temp_dir.name, "test_persistent_queue.json")
|
||||
self.patch_file = patch("nodes.queue_control.PERSISTENT_FILE", self.temp_file)
|
||||
self.patch_user = patch("nodes.queue_control.USER_DIR", self.temp_dir.name)
|
||||
self.patch_enabled = patch("nodes.queue_control.is_persistent_queue_enabled", return_value=True)
|
||||
self.patch_file.start()
|
||||
self.patch_user.start()
|
||||
self.patch_enabled.start()
|
||||
|
||||
def tearDown(self):
|
||||
self.patch_enabled.stop()
|
||||
self.patch_user.stop()
|
||||
self.patch_file.stop()
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
def test_finished_prompt_cannot_be_resynced(self):
|
||||
mgr = PersistentQueueManager()
|
||||
mgr.set_batch_info("p1", {"batchId": "b1"})
|
||||
mgr.add_item((0, "p1", {}, {}, {}))
|
||||
mgr.remove_item("p1")
|
||||
|
||||
# A late /saturnnodes/batch_queue/sync from another client must not
|
||||
# resurrect metadata for a prompt that already finished.
|
||||
mgr.set_batch_info("p1", {"batchId": "b1"})
|
||||
self.assertNotIn("p1", mgr.batch_meta)
|
||||
|
||||
# Flush so the debounce timer cannot fire after tearDown restores the
|
||||
# real PERSISTENT_FILE (it would write into the live user directory).
|
||||
mgr.flush()
|
||||
|
||||
def test_batch_meta_is_capped_keeping_recent_entries(self):
|
||||
mgr = PersistentQueueManager()
|
||||
mgr.MAX_BATCH_META = 3
|
||||
with patch("nodes.queue_control.time.time", side_effect=[1, 2, 3, 4, 5]):
|
||||
for i in range(5):
|
||||
mgr.set_batch_info(f"p{i}", {"batchId": f"b{i}"})
|
||||
|
||||
self.assertEqual(len(mgr.batch_meta), 3)
|
||||
self.assertNotIn("p0", mgr.batch_meta)
|
||||
self.assertNotIn("p1", mgr.batch_meta)
|
||||
self.assertIn("p2", mgr.batch_meta)
|
||||
self.assertIn("p4", mgr.batch_meta)
|
||||
|
||||
def test_saves_are_debounced_compact_and_flushable(self):
|
||||
mgr = PersistentQueueManager()
|
||||
mgr.add_item((0, "p1", {"prompt": {}}, {}, {}))
|
||||
# The first write is deliberately deferred; memory is authoritative.
|
||||
self.assertTrue(mgr._save_dirty)
|
||||
mgr.flush()
|
||||
|
||||
self.assertFalse(mgr._save_dirty)
|
||||
self.assertTrue(os.path.exists(self.temp_file))
|
||||
with open(self.temp_file, "r", encoding="utf-8") as f:
|
||||
raw = f.read()
|
||||
# Compact JSON (no indentation/newlines) keeps the file small.
|
||||
self.assertNotIn("\n", raw)
|
||||
payload = json.loads(raw)
|
||||
self.assertEqual(payload["schema_version"], queue_control.PERSISTENT_QUEUE_SCHEMA_VERSION)
|
||||
self.assertEqual(payload["items"][0]["prompt_id"], "p1")
|
||||
self.assertFalse(os.path.exists(self.temp_file + ".tmp"))
|
||||
|
||||
def test_restore_retry_setting_keeps_then_drops_failed_prompts(self):
|
||||
mgr = PersistentQueueManager()
|
||||
mgr.persistent_items = [{
|
||||
"prompt_id": "p1",
|
||||
"item": [0, "p1", {"prompt": {}}, {}],
|
||||
"retries": 0,
|
||||
}]
|
||||
|
||||
server = MagicMock()
|
||||
server.prompt_queue.get_history.return_value = {}
|
||||
server.prompt_queue.put.side_effect = RuntimeError("queue unavailable")
|
||||
server.number = 0
|
||||
|
||||
def fake_setting(key, default):
|
||||
return "1" if key == "PERSISTENT_QUEUE_MAX_RETRIES" else default
|
||||
|
||||
with patch("nodes.queue_control.PromptServer.instance", server), \
|
||||
patch("nodes.queue_control.get_env_setting", side_effect=fake_setting):
|
||||
mgr.restore_queue("test-client")
|
||||
self.assertEqual(len(mgr.persistent_items), 1)
|
||||
self.assertEqual(mgr.persistent_items[0]["retries"], 1)
|
||||
|
||||
mgr.has_claimed_once = False
|
||||
mgr.restore_queue("test-client")
|
||||
self.assertEqual(len(mgr.persistent_items), 0)
|
||||
self.assertIn("p1", mgr.recently_finished)
|
||||
|
||||
def test_pause_state_persists_when_enabled(self):
|
||||
def fake_setting(key, default):
|
||||
if key == "PERSIST_PAUSE_STATE":
|
||||
return "true"
|
||||
return default
|
||||
|
||||
with patch("nodes.queue_control.get_env_setting", side_effect=fake_setting):
|
||||
first = PauseQueueManager()
|
||||
first.set_pause(True, "instantly")
|
||||
second = PauseQueueManager()
|
||||
|
||||
self.assertTrue(second.paused)
|
||||
self.assertEqual(second.mode, "instantly")
|
||||
self.assertEqual(second.is_waiting, True)
|
||||
|
||||
|
||||
class TestPauseQueuePatchedGet(unittest.TestCase):
|
||||
def test_finite_timeout_returns_instead_of_busy_spinning(self):
|
||||
class FakeQueue:
|
||||
def __init__(self):
|
||||
self.currently_running = {}
|
||||
self.mutex = threading.RLock()
|
||||
self.calls = 0
|
||||
self.last_timeout = object()
|
||||
|
||||
def get(self, timeout=None):
|
||||
self.calls += 1
|
||||
self.last_timeout = timeout
|
||||
return None
|
||||
|
||||
queue = FakeQueue()
|
||||
server = MagicMock()
|
||||
server.prompt_queue = queue
|
||||
|
||||
def fake_setting(key, default):
|
||||
return default
|
||||
|
||||
with patch("nodes.queue_control.get_env_setting", side_effect=fake_setting), \
|
||||
patch.object(queue_control.PromptServer, "instance", server):
|
||||
manager = PauseQueueManager()
|
||||
manager.patch_all()
|
||||
# ComfyUI main.py calls get(timeout=0.0) when a GC pass is due.
|
||||
# The wrapper must return immediately (not spin forever).
|
||||
self.assertIsNone(queue.get(timeout=0.0))
|
||||
self.assertEqual(queue.calls, 1)
|
||||
self.assertEqual(queue.last_timeout, 0.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,363 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "batch_queue.js");
|
||||
|
||||
class FakeElement {
|
||||
constructor(tagName) {
|
||||
this.tagName = String(tagName).toUpperCase();
|
||||
this.nodeType = 1;
|
||||
this.attributes = new Map();
|
||||
this.children = [];
|
||||
this.parentNode = null;
|
||||
this.className = "";
|
||||
this.id = "";
|
||||
this.innerHTML = "";
|
||||
this.style = {};
|
||||
this.dataset = {};
|
||||
}
|
||||
|
||||
setAttribute(name, value) {
|
||||
this.attributes.set(String(name), String(value));
|
||||
}
|
||||
|
||||
getAttribute(name) {
|
||||
return this.attributes.has(String(name)) ? this.attributes.get(String(name)) : null;
|
||||
}
|
||||
|
||||
appendChild(child) {
|
||||
child.parentNode = this;
|
||||
this.children.push(child);
|
||||
return child;
|
||||
}
|
||||
|
||||
remove() {
|
||||
if (!this.parentNode) return;
|
||||
const index = this.parentNode.children.indexOf(this);
|
||||
if (index >= 0) this.parentNode.children.splice(index, 1);
|
||||
this.parentNode = null;
|
||||
}
|
||||
|
||||
querySelector(selector) {
|
||||
if (selector === ".saturnnodes-batch-line-wrap") {
|
||||
return this.children.find(child => child.className === "saturnnodes-batch-line-wrap") || null;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
matches(selector) {
|
||||
return String(selector).split(",").some(rawPart => {
|
||||
const part = rawPart.trim();
|
||||
if (part === "[data-job-id]") return this.attributes.has("data-job-id");
|
||||
if (part === '[data-testid="job-assets-list"]') {
|
||||
return this.getAttribute("data-testid") === "job-assets-list";
|
||||
}
|
||||
if (part === ".comfy-table") {
|
||||
return this.className.split(/\s+/).includes("comfy-table");
|
||||
}
|
||||
if (part === "tr") return this.tagName === "TR";
|
||||
if (part === ".comfy-table tr[data-job-id]") {
|
||||
return this.tagName === "TR" && this.attributes.has("data-job-id");
|
||||
}
|
||||
return false;
|
||||
});
|
||||
}
|
||||
|
||||
closest(selector) {
|
||||
let node = this;
|
||||
while (node) {
|
||||
if (node.matches?.(selector)) return node;
|
||||
node = node.parentNode;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
class FakeMutationObserver {
|
||||
constructor(callback) {
|
||||
this.callback = callback;
|
||||
this.target = null;
|
||||
this.options = null;
|
||||
FakeMutationObserver.instances.push(this);
|
||||
}
|
||||
|
||||
observe(target, options) {
|
||||
this.target = target;
|
||||
this.options = options;
|
||||
}
|
||||
}
|
||||
FakeMutationObserver.instances = [];
|
||||
|
||||
let clock = 1_000_000;
|
||||
let timerId = 0;
|
||||
const timers = [];
|
||||
|
||||
function fakeSetTimeout(callback, delay = 0) {
|
||||
const timer = {
|
||||
id: ++timerId,
|
||||
at: clock + Math.max(0, Number(delay) || 0),
|
||||
callback,
|
||||
};
|
||||
timers.push(timer);
|
||||
return timer.id;
|
||||
}
|
||||
|
||||
function fakeClearTimeout(id) {
|
||||
const index = timers.findIndex(timer => timer.id === id);
|
||||
if (index >= 0) timers.splice(index, 1);
|
||||
}
|
||||
|
||||
function advance(ms) {
|
||||
const end = clock + Math.max(0, Number(ms) || 0);
|
||||
while (true) {
|
||||
let nextIndex = -1;
|
||||
for (let i = 0; i < timers.length; i++) {
|
||||
if (timers[i].at <= end && (nextIndex < 0 || timers[i].at < timers[nextIndex].at)) {
|
||||
nextIndex = i;
|
||||
}
|
||||
}
|
||||
if (nextIndex < 0) break;
|
||||
const [timer] = timers.splice(nextIndex, 1);
|
||||
clock = Math.max(clock, timer.at);
|
||||
timer.callback();
|
||||
}
|
||||
clock = end;
|
||||
}
|
||||
|
||||
const queueRoot = new FakeElement("div");
|
||||
queueRoot.setAttribute("data-testid", "job-assets-list");
|
||||
const activeCard = new FakeElement("div");
|
||||
activeCard.setAttribute("data-job-id", "pid-a");
|
||||
const pendingCard = new FakeElement("div");
|
||||
pendingCard.setAttribute("data-job-id", "pid-b");
|
||||
queueRoot.appendChild(activeCard);
|
||||
queueRoot.appendChild(pendingCard);
|
||||
|
||||
// Legacy v1 rows are inserted separately by the test.
|
||||
const v1Rows = [];
|
||||
|
||||
const document = {
|
||||
body: new FakeElement("body"),
|
||||
head: new FakeElement("head"),
|
||||
getElementById(id) {
|
||||
const find = node => {
|
||||
if (node.id === id) return node;
|
||||
for (const child of node.children) {
|
||||
const found = find(child);
|
||||
if (found) return found;
|
||||
}
|
||||
return null;
|
||||
};
|
||||
return find(this.body) || find(this.head);
|
||||
},
|
||||
createElement(tagName) {
|
||||
return new FakeElement(tagName);
|
||||
},
|
||||
querySelectorAll(selector) {
|
||||
if (selector.includes(".comfy-table")) {
|
||||
return [activeCard, pendingCard, ...v1Rows];
|
||||
}
|
||||
if (selector.includes("[data-job-id]")) {
|
||||
return [activeCard, pendingCard];
|
||||
}
|
||||
if (selector === ".saturnnodes-batch-line-wrap") {
|
||||
const wraps = [];
|
||||
const find = node => {
|
||||
for (const child of node.children) {
|
||||
if (child.className === "saturnnodes-batch-line-wrap") wraps.push(child);
|
||||
find(child);
|
||||
}
|
||||
};
|
||||
find(this.body);
|
||||
return wraps;
|
||||
}
|
||||
return [];
|
||||
},
|
||||
};
|
||||
document.body.appendChild(queueRoot);
|
||||
|
||||
const storage = new Map();
|
||||
const localStorage = {
|
||||
getItem(key) {
|
||||
return storage.has(key) ? storage.get(key) : null;
|
||||
},
|
||||
setItem(key, value) {
|
||||
storage.set(key, String(value));
|
||||
},
|
||||
};
|
||||
|
||||
function batchInfo(promptId, batchId, color, itemIndex, batchCount = 2) {
|
||||
return {
|
||||
prompt_id: promptId,
|
||||
batchId,
|
||||
batchNumber: 1,
|
||||
batchCount,
|
||||
itemIndex,
|
||||
color,
|
||||
colorIndex: 0,
|
||||
queuedAt: clock,
|
||||
};
|
||||
}
|
||||
|
||||
const oldBatch = "batch_old";
|
||||
const newBatch = "batch_new";
|
||||
storage.set(
|
||||
"saturnnodes_batch_queue_meta",
|
||||
JSON.stringify([
|
||||
batchInfo("pid-a", oldBatch, "#4ade80", 0, 1),
|
||||
batchInfo("pid-b", newBatch, "#fbbf24", 1, 2),
|
||||
batchInfo("pid-c", newBatch, "#fbbf24", 0, 2),
|
||||
batchInfo("pid-d", newBatch, "#fbbf24", 1, 2),
|
||||
])
|
||||
);
|
||||
|
||||
const extension = {
|
||||
name: "test",
|
||||
setup: null,
|
||||
};
|
||||
const app = {
|
||||
queuePrompt: async () => true,
|
||||
extensionManager: {
|
||||
setting: {
|
||||
get: () => true,
|
||||
},
|
||||
},
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.setup = candidate.setup;
|
||||
},
|
||||
};
|
||||
const apiEvents = new Map();
|
||||
const api = {
|
||||
queuePrompt: async () => ({ prompt_id: "queued" }),
|
||||
fetchApi: async () => null,
|
||||
addEventListener(name, callback) {
|
||||
if (!apiEvents.has(name)) apiEvents.set(name, []);
|
||||
apiEvents.get(name).push(callback);
|
||||
},
|
||||
};
|
||||
|
||||
const source = fs
|
||||
.readFileSync(sourcePath, "utf8")
|
||||
.replace(/^import .*$/gm, "")
|
||||
.replace(/export \{[^}]*\};?\s*$/m, "");
|
||||
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
app,
|
||||
api,
|
||||
authenticatedFetch: async () => null,
|
||||
document,
|
||||
localStorage,
|
||||
MutationObserver: FakeMutationObserver,
|
||||
setTimeout: fakeSetTimeout,
|
||||
clearTimeout: fakeClearTimeout,
|
||||
Date: { now: () => clock },
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
});
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
|
||||
assert.equal(typeof extension.setup, "function", "batch queue extension must register");
|
||||
await extension.setup();
|
||||
advance(100);
|
||||
|
||||
const activeLine = activeCard.querySelector(".saturnnodes-batch-line-wrap");
|
||||
assert.ok(activeLine, "active card must have a batch line");
|
||||
assert.match(activeLine.innerHTML, /#4ade80/, "initial active line must use its own batch color");
|
||||
|
||||
const observer = FakeMutationObserver.instances.find(
|
||||
item => item.options && item.options.attributes && item.options.attributeFilter?.includes("data-job-id")
|
||||
);
|
||||
assert.ok(observer, "queue observer must watch data-job-id attribute changes");
|
||||
|
||||
// Simulate Vue reusing the running card for the first item of the next batch.
|
||||
// The old green SVG remains in the DOM until the attribute mutation is handled.
|
||||
activeCard.setAttribute("data-job-id", "pid-c");
|
||||
observer.callback([{ target: activeCard, type: "attributes", attributeName: "data-job-id" }]);
|
||||
advance(100);
|
||||
|
||||
assert.match(
|
||||
activeCard.querySelector(".saturnnodes-batch-line-wrap").innerHTML,
|
||||
/#fbbf24/,
|
||||
"reused active card must repaint with the new batch color"
|
||||
);
|
||||
assert.doesNotMatch(
|
||||
activeCard.querySelector(".saturnnodes-batch-line-wrap").innerHTML,
|
||||
/#4ade80/,
|
||||
"stale batch color must not survive a data-job-id change"
|
||||
);
|
||||
|
||||
// The new batch's pending item and active item must share the same color.
|
||||
const pendingLine = pendingCard.querySelector(".saturnnodes-batch-line-wrap");
|
||||
assert.ok(pendingLine, "pending card must have a batch line");
|
||||
assert.match(pendingLine.innerHTML, /#fbbf24/, "pending item must use the new batch color");
|
||||
|
||||
// 4. Legacy v1: a table can be inserted with its rows already built, so the
|
||||
// mutation target is the new parent and the table only appears in addedNodes.
|
||||
const table = new FakeElement("table");
|
||||
table.className = "comfy-table";
|
||||
const v1Row = new FakeElement("tr");
|
||||
v1Row.setAttribute("data-job-id", "pid-d");
|
||||
table.appendChild(v1Row);
|
||||
const tableWrapper = new FakeElement("div");
|
||||
document.body.appendChild(tableWrapper);
|
||||
tableWrapper.appendChild(table);
|
||||
v1Rows.push(v1Row);
|
||||
|
||||
observer.callback([{
|
||||
type: "childList",
|
||||
target: tableWrapper,
|
||||
addedNodes: [table],
|
||||
removedNodes: [],
|
||||
}]);
|
||||
advance(100);
|
||||
|
||||
const v1Line = v1Row.querySelector(".saturnnodes-batch-line-wrap");
|
||||
assert.ok(v1Line, "a v1 table inserted with rows must get a batch line");
|
||||
assert.match(v1Line.innerHTML, /#fbbf24/, "v1 row must use the batch color");
|
||||
|
||||
// 5. Duplicate data-job-id rows must not skew batch grouping or paint twice.
|
||||
const duplicateRow = new FakeElement("tr");
|
||||
duplicateRow.setAttribute("data-job-id", "pid-b");
|
||||
// Simulate a line left over from an earlier render on the duplicate.
|
||||
const staleDuplicateLine = new FakeElement("div");
|
||||
staleDuplicateLine.className = "saturnnodes-batch-line-wrap";
|
||||
duplicateRow.appendChild(staleDuplicateLine);
|
||||
table.appendChild(duplicateRow);
|
||||
v1Rows.push(duplicateRow);
|
||||
|
||||
observer.callback([{
|
||||
type: "childList",
|
||||
target: table,
|
||||
addedNodes: [duplicateRow],
|
||||
removedNodes: [],
|
||||
}]);
|
||||
advance(100);
|
||||
|
||||
assert.equal(
|
||||
duplicateRow.querySelector(".saturnnodes-batch-line-wrap"),
|
||||
null,
|
||||
"duplicate job row must be ignored"
|
||||
);
|
||||
assert.ok(
|
||||
pendingCard.querySelector(".saturnnodes-batch-line-wrap"),
|
||||
"the first occurrence of the job id keeps its batch line"
|
||||
);
|
||||
|
||||
console.log("batch queue DOM regression test passed");
|
||||
@@ -0,0 +1,30 @@
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import unittest
|
||||
|
||||
|
||||
class TestBatchQueueDomRegression(unittest.TestCase):
|
||||
def test_reused_active_card_repaints_after_job_id_change(self):
|
||||
node = shutil.which("node")
|
||||
if not node:
|
||||
self.skipTest("Node.js is not available")
|
||||
|
||||
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
script = os.path.join(repo_root, "tests", "test_batch_queue_dom.mjs")
|
||||
result = subprocess.run(
|
||||
[node, script],
|
||||
cwd=repo_root,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
self.assertEqual(
|
||||
result.returncode,
|
||||
0,
|
||||
msg=f"Batch queue DOM regression failed:\n{result.stdout}\n{result.stderr}",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,339 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { FakeMutationObserver, createDocument, stripImports } from "./fake_dom.mjs";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "batch_queue.js");
|
||||
|
||||
let source = stripImports(fs.readFileSync(sourcePath, "utf8"));
|
||||
source += `
|
||||
globalThis.__batchTest = {
|
||||
updateQueueBatchVisuals,
|
||||
normalizeBatchColor,
|
||||
normalizeSegmentType,
|
||||
buildSvgPath,
|
||||
getBatchRegistry: () => batchRegistry,
|
||||
};
|
||||
`;
|
||||
|
||||
let clock = 1_000_000;
|
||||
let timerId = 0;
|
||||
const timers = [];
|
||||
|
||||
function fakeSetTimeout(callback, delay = 0) {
|
||||
const timer = {
|
||||
id: ++timerId,
|
||||
at: clock + Math.max(0, Number(delay) || 0),
|
||||
callback,
|
||||
};
|
||||
timers.push(timer);
|
||||
return timer.id;
|
||||
}
|
||||
|
||||
function fakeClearTimeout(id) {
|
||||
const index = timers.findIndex(timer => timer.id === id);
|
||||
if (index >= 0) timers.splice(index, 1);
|
||||
}
|
||||
|
||||
function advance(ms) {
|
||||
const end = clock + Math.max(0, Number(ms) || 0);
|
||||
while (true) {
|
||||
let nextIndex = -1;
|
||||
for (let i = 0; i < timers.length; i++) {
|
||||
if (timers[i].at <= end && (nextIndex < 0 || timers[i].at < timers[nextIndex].at)) {
|
||||
nextIndex = i;
|
||||
}
|
||||
}
|
||||
if (nextIndex < 0) break;
|
||||
const [timer] = timers.splice(nextIndex, 1);
|
||||
clock = Math.max(clock, timer.at);
|
||||
timer.callback();
|
||||
}
|
||||
clock = end;
|
||||
}
|
||||
|
||||
const document = createDocument();
|
||||
const queueRoot = document.createElement("div");
|
||||
queueRoot.setAttribute("data-testid", "job-assets-list");
|
||||
document.body.appendChild(queueRoot);
|
||||
|
||||
const storage = new Map();
|
||||
const localStorage = {
|
||||
getItem(key) {
|
||||
return storage.has(key) ? storage.get(key) : null;
|
||||
},
|
||||
setItem(key, value) {
|
||||
storage.set(key, String(value));
|
||||
},
|
||||
};
|
||||
|
||||
function info(promptId, batchId, color, itemIndex, batchCount, queuedAt) {
|
||||
return {
|
||||
prompt_id: promptId,
|
||||
batchId,
|
||||
batchNumber: 1,
|
||||
batchCount,
|
||||
itemIndex,
|
||||
color,
|
||||
colorIndex: 0,
|
||||
queuedAt,
|
||||
};
|
||||
}
|
||||
|
||||
const seeded = [];
|
||||
for (let index = 0; index < 260; index++) {
|
||||
seeded.push(info(`old-${index}`, `old-batch-${index}`, "#10b981", 0, 1, index + 1));
|
||||
}
|
||||
seeded.push(info("evil-a", "batch-evil", '#fff" onload="alert(1)', 0, 2, 1000));
|
||||
seeded.push(info("evil-b", "batch-evil", "</style><script>alert(1)</script>", 1, 2, 1001));
|
||||
seeded.push(info("bad-count", "batch-bad", "#123456", "not-an-int", -5, 1002));
|
||||
seeded.push(info("evil-c", "batch-evil-c", '#fff" onload="alert(1)', 0, 1, 1003));
|
||||
seeded.push(info("low-index", "batch-low-index", "#123456", -5, 3, 1004));
|
||||
seeded.push(info("high-index", "batch-high-index", "#123456", 99, 3, 1005));
|
||||
storage.set("saturnnodes_batch_queue_meta", JSON.stringify(seeded));
|
||||
storage.set("saturnnodes_batch_counter", "-12");
|
||||
|
||||
const serverBatchData = {
|
||||
// Older server data must not replace newer local metadata.
|
||||
"old-0": info("old-0", "old-batch-0", "#10b981", 0, 1, 0),
|
||||
// Newer server data must replace older local metadata.
|
||||
"old-1": info("old-1", "old-batch-1", "#10b981", 0, 1, 9999),
|
||||
"server-only": info("server-only", "server-batch", "#3b82f6", 0, 1, 2000),
|
||||
};
|
||||
|
||||
const apiEvents = new Map();
|
||||
const deferred = [];
|
||||
let nextPid = 0;
|
||||
const api = {
|
||||
clientId: "client-test",
|
||||
queuePrompt: () => new Promise(resolve => {
|
||||
deferred.push(() => resolve({ prompt_id: `conc-${++nextPid}` }));
|
||||
}),
|
||||
fetchApi: async url => {
|
||||
if (url === "/saturnnodes/batch_queue/data") {
|
||||
return { ok: true, json: async () => serverBatchData };
|
||||
}
|
||||
return null;
|
||||
},
|
||||
addEventListener(name, callback) {
|
||||
if (!apiEvents.has(name)) apiEvents.set(name, []);
|
||||
apiEvents.get(name).push(callback);
|
||||
},
|
||||
};
|
||||
|
||||
const extension = { name: "", setup: null };
|
||||
const app = {
|
||||
extensionManager: { setting: { get: () => true } },
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.setup = candidate.setup;
|
||||
},
|
||||
async queuePrompt(number, batchCount, queueNodeIds) {
|
||||
// Both submissions start before either response arrives. This is the
|
||||
// concurrency window that used to race on prompts.length.
|
||||
const first = api.queuePrompt(number, { output: {} }, queueNodeIds);
|
||||
const second = api.queuePrompt(number, { output: {} }, queueNodeIds);
|
||||
return await Promise.all([first, second]);
|
||||
},
|
||||
};
|
||||
|
||||
const mathStub = Object.create(Math);
|
||||
let randomSeed = 0;
|
||||
mathStub.random = () => ((++randomSeed * 0.6180339887) % 1);
|
||||
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
app,
|
||||
api,
|
||||
authenticatedFetch: async () => ({ ok: true }),
|
||||
document,
|
||||
localStorage,
|
||||
MutationObserver: FakeMutationObserver,
|
||||
setTimeout: fakeSetTimeout,
|
||||
clearTimeout: fakeClearTimeout,
|
||||
Date: { now: () => clock },
|
||||
Math: mathStub,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
});
|
||||
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
const batchTest = context.__batchTest;
|
||||
|
||||
assert.equal(typeof extension.setup, "function", "batch queue extension must register");
|
||||
await extension.setup();
|
||||
await new Promise(resolve => setImmediate(resolve));
|
||||
advance(100);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Restored storage is pruned by queuedAt, not by Map insertion order.
|
||||
// ---------------------------------------------------------------------------
|
||||
const savedAfterLoad = JSON.parse(storage.get("saturnnodes_batch_queue_meta"));
|
||||
assert.equal(savedAfterLoad.length, 250, "restored metadata must be pruned to 250 entries");
|
||||
const savedIds = new Set(savedAfterLoad.map(entry => entry.prompt_id));
|
||||
assert.equal(savedIds.has("old-0"), false, "oldest restored metadata must be pruned");
|
||||
assert.equal(savedIds.has("old-259"), true, "newest restored metadata must be kept");
|
||||
assert.equal(savedIds.has("evil-a"), true, "malicious-color entries must not be dropped");
|
||||
assert.equal(savedIds.has("server-only"), true, "newer server metadata must be merged");
|
||||
|
||||
const registry = batchTest.getBatchRegistry();
|
||||
assert.equal(registry.get("old-0").queuedAt, 1, "older server snapshot must not overwrite local metadata");
|
||||
assert.equal(registry.get("old-1").queuedAt, 9999, "newer server snapshot must replace local metadata");
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Restored rows with injected color strings must render safe hex colors only.
|
||||
// ---------------------------------------------------------------------------
|
||||
function makeRow(promptId) {
|
||||
const row = document.createElement("div");
|
||||
row.setAttribute("data-job-id", promptId);
|
||||
queueRoot.appendChild(row);
|
||||
return row;
|
||||
}
|
||||
|
||||
const evilA = makeRow("evil-a");
|
||||
const evilB = makeRow("evil-b");
|
||||
const badCount = makeRow("bad-count");
|
||||
const lowIndex = makeRow("low-index");
|
||||
const highIndex = makeRow("high-index");
|
||||
const duplicateEvilA = makeRow("evil-a");
|
||||
|
||||
await batchTest.updateQueueBatchVisuals();
|
||||
|
||||
const evilALine = evilA.querySelector(".saturnnodes-batch-line-wrap");
|
||||
const evilBLine = evilB.querySelector(".saturnnodes-batch-line-wrap");
|
||||
assert.ok(evilALine, "restored row must receive a batch line");
|
||||
assert.ok(evilBLine, "restored row must receive a batch line");
|
||||
assert.match(evilALine.innerHTML, /#10b981/, "invalid color must fall back to the default hex color");
|
||||
assert.doesNotMatch(evilALine.innerHTML, /alert|onload|<script/i, "color injection must never reach innerHTML");
|
||||
assert.match(evilBLine.innerHTML, /#10b981/, "markup color must fall back to the default hex color");
|
||||
assert.doesNotMatch(evilBLine.innerHTML, /alert|<script|<\/style/i, "markup color must never reach innerHTML");
|
||||
assert.match(
|
||||
evilALine.innerHTML,
|
||||
/M 9,6 Q 2,6 2,16 L 2,48/,
|
||||
"the first visible item of a restored batch must draw the first segment"
|
||||
);
|
||||
assert.match(
|
||||
evilBLine.innerHTML,
|
||||
/M 2,0 L 2,32 Q 2,42 9,42/,
|
||||
"the last visible item of a restored batch must draw the last segment"
|
||||
);
|
||||
assert.equal(
|
||||
duplicateEvilA.querySelector(".saturnnodes-batch-line-wrap"),
|
||||
null,
|
||||
"a duplicate restored row must not receive a second line"
|
||||
);
|
||||
|
||||
const badLine = badCount.querySelector(".saturnnodes-batch-line-wrap");
|
||||
assert.ok(badLine, "invalid batch counters must still render a line");
|
||||
assert.match(badLine.innerHTML, /#123456/, "valid custom hex colors must be preserved");
|
||||
assert.match(
|
||||
badLine.innerHTML,
|
||||
/M 9,6 Q 2,6 2,16 L 2,32 Q 2,42 9,42/,
|
||||
"negative batchCount and non-integer itemIndex must clamp to a single segment"
|
||||
);
|
||||
|
||||
const lowIndexLine = lowIndex.querySelector(".saturnnodes-batch-line-wrap");
|
||||
assert.ok(lowIndexLine, "negative itemIndex must still render a line");
|
||||
assert.match(
|
||||
lowIndexLine.innerHTML,
|
||||
/M 9,6 Q 2,6 2,16 L 2,48/,
|
||||
"a negative itemIndex must clamp to the first segment"
|
||||
);
|
||||
|
||||
const highIndexLine = highIndex.querySelector(".saturnnodes-batch-line-wrap");
|
||||
assert.ok(highIndexLine, "out-of-range itemIndex must still render a line");
|
||||
assert.match(
|
||||
highIndexLine.innerHTML,
|
||||
/M 2,0 L 2,32 Q 2,42 9,42/,
|
||||
"an out-of-range itemIndex must clamp to the last segment"
|
||||
);
|
||||
|
||||
// Direct sanitizer checks.
|
||||
assert.equal(batchTest.normalizeBatchColor("#ABC"), "#ABC");
|
||||
assert.equal(batchTest.normalizeBatchColor("#ABCD"), "#ABCD");
|
||||
assert.equal(batchTest.normalizeBatchColor("#123456"), "#123456");
|
||||
assert.equal(batchTest.normalizeBatchColor("#12345678"), "#12345678");
|
||||
assert.equal(batchTest.normalizeBatchColor(" #10b981 "), "#10b981");
|
||||
assert.equal(batchTest.normalizeBatchColor("#12345"), "#10b981");
|
||||
assert.equal(batchTest.normalizeBatchColor("#1234567"), "#10b981");
|
||||
assert.equal(batchTest.normalizeBatchColor("red"), "#10b981");
|
||||
assert.equal(batchTest.normalizeBatchColor("url(javascript:alert(1))"), "#10b981");
|
||||
assert.equal(batchTest.normalizeSegmentType("bogus"), "single");
|
||||
assert.equal(batchTest.normalizeSegmentType('"><script>alert(1)</script>'), "single");
|
||||
assert.equal(batchTest.buildSvgPath("bogus"), batchTest.buildSvgPath("single"));
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// A negative persisted counter must clamp to a valid BATCH_COLORS index.
|
||||
// ---------------------------------------------------------------------------
|
||||
const submission = app.queuePrompt(-1, 2, null);
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
assert.equal(deferred.length, 2, "both concurrent submissions must reserve an item index");
|
||||
deferred.shift()();
|
||||
deferred.shift()();
|
||||
await submission;
|
||||
|
||||
const afterBatch = JSON.parse(storage.get("saturnnodes_batch_queue_meta"));
|
||||
const concOne = afterBatch.find(entry => entry.prompt_id === "conc-1");
|
||||
const concTwo = afterBatch.find(entry => entry.prompt_id === "conc-2");
|
||||
assert.ok(concOne && concTwo, "both concurrent prompts must be recorded");
|
||||
assert.equal(concOne.batchId, concTwo.batchId, "concurrent prompts must share one batch");
|
||||
assert.deepEqual(
|
||||
[concOne.itemIndex, concTwo.itemIndex].sort((a, b) => a - b),
|
||||
[0, 1],
|
||||
"concurrent submissions must reserve unique item indexes"
|
||||
);
|
||||
assert.equal(concOne.color, "#10b981", "a negative persisted counter must clamp to color index 0");
|
||||
assert.equal(concOne.colorIndex, 0);
|
||||
assert.equal(concTwo.batchCount, 2);
|
||||
|
||||
// The completed batch must release the active context, so the next standalone
|
||||
// submission starts a fresh single-item batch.
|
||||
const standalone = api.queuePrompt(0, { output: {} }, null);
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
assert.equal(deferred.length, 1, "standalone submission must reach the queue hook");
|
||||
deferred.shift()();
|
||||
await standalone;
|
||||
|
||||
const afterStandalone = JSON.parse(storage.get("saturnnodes_batch_queue_meta"));
|
||||
const concThree = afterStandalone.find(entry => entry.prompt_id === "conc-3");
|
||||
assert.ok(concThree, "standalone prompt must be recorded");
|
||||
assert.notEqual(concThree.batchId, concOne.batchId, "completed batch context must be cleared");
|
||||
assert.equal(concThree.batchCount, 1);
|
||||
assert.equal(concThree.color, "#3b82f6", "the next batch must advance the color counter");
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Reused restored cards repaint after Vue changes data-job-id.
|
||||
// ---------------------------------------------------------------------------
|
||||
const observer = FakeMutationObserver.instances.find(
|
||||
item => item.target === document.body && item.options?.attributes
|
||||
);
|
||||
assert.ok(observer, "queue observer must watch attribute changes");
|
||||
const reused = makeRow("old-259");
|
||||
await batchTest.updateQueueBatchVisuals();
|
||||
assert.ok(reused.querySelector(".saturnnodes-batch-line-wrap"), "reused card must get a line");
|
||||
|
||||
reused.setAttribute("data-job-id", "evil-c");
|
||||
observer.emit([{ target: reused, type: "attributes", attributeName: "data-job-id" }]);
|
||||
advance(100);
|
||||
assert.match(
|
||||
reused.querySelector(".saturnnodes-batch-line-wrap").innerHTML,
|
||||
/#10b981/,
|
||||
"reused card must repaint with the new job's safe batch color"
|
||||
);
|
||||
|
||||
console.log("batch queue restore DOM regression test passed");
|
||||
@@ -0,0 +1,309 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
|
||||
// Regression coverage for the queue-path bleed reported in the LoRA loader
|
||||
// audit follow-up:
|
||||
//
|
||||
// Queue prompt 1, then queue prompt 2 while prompt 1 is still being
|
||||
// submitted, then edit the canvas. ComfyUI serializes prompt 2 only after
|
||||
// prompt 1's submission turn finishes, so without a click-time snapshot the
|
||||
// edited text/LoRA selection bleeds into prompt 2.
|
||||
//
|
||||
// The fake ComfyUI queue below mirrors the real `app.queuePrompt` loop: it
|
||||
// pushes items, returns early while `processingQueue` is true, and calls
|
||||
// `graphToPrompt()` freshly for every submitted item.
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "batch_queue.js");
|
||||
|
||||
class FakeStyleElement {
|
||||
constructor() {
|
||||
this.textContent = "";
|
||||
this.id = "";
|
||||
this.children = [];
|
||||
}
|
||||
remove() {}
|
||||
appendChild(child) {
|
||||
this.children.push(child);
|
||||
return child;
|
||||
}
|
||||
querySelector() {
|
||||
return null;
|
||||
}
|
||||
querySelectorAll() {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
const document = {
|
||||
body: new FakeStyleElement(),
|
||||
head: new FakeStyleElement(),
|
||||
getElementById() {
|
||||
return null;
|
||||
},
|
||||
createElement() {
|
||||
return new FakeStyleElement();
|
||||
},
|
||||
querySelector() {
|
||||
return null;
|
||||
},
|
||||
querySelectorAll() {
|
||||
return [];
|
||||
},
|
||||
};
|
||||
|
||||
const storage = new Map();
|
||||
const localStorage = {
|
||||
getItem(key) {
|
||||
return storage.has(key) ? storage.get(key) : null;
|
||||
},
|
||||
setItem(key, value) {
|
||||
storage.set(key, String(value));
|
||||
},
|
||||
};
|
||||
|
||||
class FakeMutationObserver {
|
||||
constructor(callback) {
|
||||
this.callback = callback;
|
||||
}
|
||||
observe() {}
|
||||
disconnect() {}
|
||||
}
|
||||
|
||||
// --- Fake ComfyUI canvas + queue ---------------------------------------------
|
||||
|
||||
const canvas = { text: "", lora: "[]", seed: 1 };
|
||||
const submissions = [];
|
||||
const pendingResponses = [];
|
||||
let graphToPromptCalls = 0;
|
||||
let nextPromptId = 0;
|
||||
|
||||
async function graphToPrompt() {
|
||||
graphToPromptCalls++;
|
||||
// Read values synchronously, like the real serializer does.
|
||||
return {
|
||||
output: {
|
||||
"10": {
|
||||
inputs: {
|
||||
text: canvas.text,
|
||||
_selected_lora: canvas.lora,
|
||||
seed: canvas.seed,
|
||||
},
|
||||
},
|
||||
},
|
||||
workflow: { nodes: [] },
|
||||
};
|
||||
}
|
||||
|
||||
async function apiQueuePrompt(number, prompt) {
|
||||
// Copy the prompt exactly as it will be sent to the server.
|
||||
submissions.push(JSON.parse(JSON.stringify(prompt)));
|
||||
await new Promise(resolve => pendingResponses.push(resolve));
|
||||
nextPromptId++;
|
||||
return { prompt_id: `pid-${nextPromptId}` };
|
||||
}
|
||||
|
||||
async function comfyQueuePrompt(number, batchCount = 1) {
|
||||
this.queueItems.push({ number, batchCount });
|
||||
if (this.processingQueue) return false;
|
||||
this.processingQueue = true;
|
||||
try {
|
||||
while (this.queueItems.length) {
|
||||
const item = this.queueItems.pop();
|
||||
for (let r = 0; r < item.batchCount; r++) {
|
||||
const prompt = await this.graphToPrompt(this.rootGraph);
|
||||
await api.queuePrompt(item.number, prompt);
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
this.processingQueue = false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
const extension = { setup: null };
|
||||
const app = {
|
||||
queuePrompt: comfyQueuePrompt,
|
||||
graphToPrompt,
|
||||
graph: { getNodeById: () => null },
|
||||
queueItems: [],
|
||||
processingQueue: false,
|
||||
extensionManager: { setting: { get: () => true } },
|
||||
registerExtension(candidate) {
|
||||
extension.setup = candidate.setup;
|
||||
},
|
||||
setDirtyCanvas() {},
|
||||
};
|
||||
|
||||
const api = {
|
||||
queuePrompt: apiQueuePrompt,
|
||||
fetchApi: async () => null,
|
||||
addEventListener() {},
|
||||
};
|
||||
|
||||
const source = fs
|
||||
.readFileSync(sourcePath, "utf8")
|
||||
.replace(/^import .*$/gm, "")
|
||||
.replace(/export \{[^}]*\};?\s*$/m, "");
|
||||
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
app,
|
||||
api,
|
||||
authenticatedFetch: async () => null,
|
||||
document,
|
||||
localStorage,
|
||||
MutationObserver: FakeMutationObserver,
|
||||
setTimeout: (callback) => setTimeout(callback, 0),
|
||||
clearTimeout,
|
||||
Date,
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
});
|
||||
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
assert.equal(typeof extension.setup, "function", "batch queue extension must register");
|
||||
await extension.setup();
|
||||
|
||||
const flush = async () => {
|
||||
for (let i = 0; i < 50; i++) await Promise.resolve();
|
||||
};
|
||||
const resolveNext = async () => {
|
||||
const resolve = pendingResponses.shift();
|
||||
assert.ok(resolve, "expected a pending API submission");
|
||||
resolve();
|
||||
await flush();
|
||||
};
|
||||
const resetCanvas = (text, lora = "[]", seed = 1) => {
|
||||
canvas.text = text;
|
||||
canvas.lora = lora;
|
||||
canvas.seed = seed;
|
||||
};
|
||||
const promptText = index => submissions[index].output["10"].inputs.text;
|
||||
const promptLora = index => submissions[index].output["10"].inputs._selected_lora;
|
||||
const promptSeed = index => submissions[index].output["10"].inputs.seed;
|
||||
|
||||
// --- Scenario 1: queued prompt must keep the text/LoRA shown at click time ---
|
||||
|
||||
resetCanvas("prompt one text", '["LoRA_A.safetensors"]');
|
||||
graphToPromptCalls = 0;
|
||||
|
||||
// First click starts submission immediately (no extra snapshot needed).
|
||||
const first = app.queuePrompt(0, 1);
|
||||
await flush();
|
||||
assert.equal(app.processingQueue, true, "first submission must be in flight");
|
||||
assert.equal(graphToPromptCalls, 1, "an idle click must not serialize twice");
|
||||
|
||||
// The user immediately clicks Queue again, then keeps editing the canvas
|
||||
// while the first prompt is still being submitted.
|
||||
const second = app.queuePrompt(0, 1);
|
||||
assert.equal(
|
||||
graphToPromptCalls,
|
||||
2,
|
||||
"a queued click during an in-flight submission must snapshot the graph at click time"
|
||||
);
|
||||
|
||||
canvas.text = "prompt two text"; // edited after both clicks
|
||||
canvas.lora = '["LoRA_B.safetensors"]';
|
||||
canvas.seed = 2;
|
||||
|
||||
await resolveNext(); // finish prompt one
|
||||
await resolveNext(); // finish prompt two
|
||||
await Promise.all([first, second]);
|
||||
await flush();
|
||||
|
||||
assert.equal(submissions.length, 2, "both queued prompts must be submitted");
|
||||
assert.equal(promptText(0), "prompt one text");
|
||||
assert.equal(promptLora(0), '["LoRA_A.safetensors"]');
|
||||
assert.equal(
|
||||
promptText(1),
|
||||
"prompt one text",
|
||||
"canvas text edited after the click must not bleed into the queued prompt"
|
||||
);
|
||||
assert.equal(
|
||||
promptLora(1),
|
||||
'["LoRA_A.safetensors"]',
|
||||
"LoRA selection edited after the click must not bleed into the queued prompt"
|
||||
);
|
||||
|
||||
// --- Scenario 2: a multi-item batch stays frozen at click time --------------
|
||||
|
||||
submissions.length = 0;
|
||||
resetCanvas("batch base text", '["LoRA_C.safetensors"]', 3);
|
||||
app.queueItems.length = 0;
|
||||
|
||||
const batch = app.queuePrompt(0, 3);
|
||||
await flush(); // item 0 in flight
|
||||
|
||||
canvas.text = "edited mid batch"; // must not leak into items 2/3
|
||||
canvas.lora = '["LoRA_D.safetensors"]';
|
||||
canvas.seed = 99;
|
||||
|
||||
await resolveNext(); // item 0
|
||||
await resolveNext(); // item 1
|
||||
await resolveNext(); // item 2
|
||||
await batch;
|
||||
await flush();
|
||||
|
||||
assert.equal(submissions.length, 3);
|
||||
for (let i = 0; i < 3; i++) {
|
||||
assert.equal(promptText(i), "batch base text", `batch item ${i + 1} text must stay frozen`);
|
||||
assert.equal(
|
||||
promptLora(i),
|
||||
'["LoRA_C.safetensors"]',
|
||||
`batch item ${i + 1} LoRA must stay frozen`
|
||||
);
|
||||
}
|
||||
|
||||
// --- Scenario 3: dynamic seeds still follow the live prompt -----------------
|
||||
|
||||
submissions.length = 0;
|
||||
resetCanvas("blocker", "[]", 10);
|
||||
const blocker = app.queuePrompt(0, 1);
|
||||
await flush(); // blocker in flight so the next click has a submission backlog
|
||||
|
||||
resetCanvas("seed batch", "[]", 11);
|
||||
const seeded = app.queuePrompt(0, 2);
|
||||
await flush();
|
||||
canvas.seed = 12; // simulates the seed widget randomizing before submission
|
||||
await resolveNext(); // blocker
|
||||
await resolveNext(); // seeded item 0
|
||||
await resolveNext(); // seeded item 1
|
||||
await blocker;
|
||||
await seeded;
|
||||
await flush();
|
||||
|
||||
assert.equal(submissions.length, 3);
|
||||
assert.equal(promptSeed(1), 12, "item 0 must use the seed the live prompt queued");
|
||||
assert.equal(promptSeed(2), 12, "item 1 must use the live seed after merging");
|
||||
assert.equal(promptText(1), "seed batch");
|
||||
assert.equal(promptText(2), "seed batch");
|
||||
|
||||
// --- Scenario 4: idle clicks keep serializing once and never double-queue ---
|
||||
|
||||
submissions.length = 0;
|
||||
resetCanvas("solo", "[]", 20);
|
||||
const solo = app.queuePrompt(0, 1);
|
||||
await flush();
|
||||
assert.equal(graphToPromptCalls > 0, true);
|
||||
await resolveNext();
|
||||
await solo;
|
||||
|
||||
assert.equal(submissions.length, 1);
|
||||
assert.equal(promptText(0), "solo");
|
||||
|
||||
console.log("batch queue click-time snapshot DOM test passed");
|
||||
@@ -0,0 +1,157 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import {
|
||||
FakeMutationObserver,
|
||||
createDocument,
|
||||
createWindow,
|
||||
stripImports,
|
||||
} from "./fake_dom.mjs";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "js", "copy_prompt.js");
|
||||
|
||||
const document = createDocument();
|
||||
const window = createWindow(document);
|
||||
|
||||
const extension = { name: "", beforeRegisterNodeDef: null };
|
||||
const app = {
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.beforeRegisterNodeDef = candidate.beforeRegisterNodeDef;
|
||||
},
|
||||
extensionManager: {
|
||||
setting: { get: () => true },
|
||||
toast: { add() {} },
|
||||
},
|
||||
ui: { settings: { get: () => undefined, getSettingValue: () => undefined } },
|
||||
graph: { setDirtyCanvas() {} },
|
||||
};
|
||||
|
||||
let source = stripImports(fs.readFileSync(sourcePath, "utf8"));
|
||||
source += `
|
||||
globalThis.__copyTest = {
|
||||
registerAssetInteraction,
|
||||
getActiveImageSrc,
|
||||
injectHoverCopyAction,
|
||||
setCopyImagePrompt: (fn) => { copyImagePrompt = fn; },
|
||||
getLastCard: () => lastInteractedAssetCard,
|
||||
};
|
||||
`;
|
||||
|
||||
const copied = [];
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
document,
|
||||
window,
|
||||
navigator: {},
|
||||
app,
|
||||
api: { addEventListener() {}, fetchApi: async () => null },
|
||||
MutationObserver: FakeMutationObserver,
|
||||
setTimeout: () => 0,
|
||||
clearTimeout: () => {},
|
||||
setInterval: () => 0,
|
||||
clearInterval: () => {},
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
Date,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
Node: { ELEMENT_NODE: 1 },
|
||||
MouseEvent: class MouseEvent {
|
||||
constructor(type, init = {}) {
|
||||
this.type = type;
|
||||
this.bubbles = !!init.bubbles;
|
||||
this.cancelable = !!init.cancelable;
|
||||
}
|
||||
},
|
||||
});
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
|
||||
assert.equal(typeof extension.beforeRegisterNodeDef, "function", "copy prompt extension must register");
|
||||
|
||||
const {
|
||||
registerAssetInteraction,
|
||||
getActiveImageSrc,
|
||||
injectHoverCopyAction,
|
||||
setCopyImagePrompt,
|
||||
getLastCard,
|
||||
} = context.__copyTest;
|
||||
|
||||
setCopyImagePrompt(async src => {
|
||||
copied.push(src);
|
||||
return true;
|
||||
});
|
||||
|
||||
function makeCard(assetId, src) {
|
||||
const card = document.createElement("div");
|
||||
card.setAttribute("data-asset-id", assetId);
|
||||
const img = document.createElement("img");
|
||||
img.src = src;
|
||||
card.appendChild(img);
|
||||
document.body.appendChild(card);
|
||||
return { card, img };
|
||||
}
|
||||
|
||||
// 1. A fresh interaction resolves the current asset.
|
||||
const first = makeCard("asset-1", "a.png");
|
||||
registerAssetInteraction(first.img);
|
||||
assert.equal(getActiveImageSrc(), "a.png", "fresh card interaction must resolve its image");
|
||||
|
||||
// 2. The virtual grid recycles the same element for another asset. The stale
|
||||
// pointer must not silently copy the new asset.
|
||||
first.card.setAttribute("data-asset-id", "asset-2");
|
||||
first.img.src = "b.png";
|
||||
assert.equal(getActiveImageSrc(), null, "recycled card must not resolve a different asset");
|
||||
assert.equal(getLastCard(), null, "recycled card must be discarded");
|
||||
|
||||
// 3. A detached card is discarded as well.
|
||||
registerAssetInteraction(first.img);
|
||||
first.card.remove();
|
||||
assert.equal(getActiveImageSrc(), null, "detached card must not resolve an image");
|
||||
assert.equal(getLastCard(), null, "detached card must be cleared");
|
||||
|
||||
// 4. Legacy hover overlay: the injected button must follow the asset that is
|
||||
// currently displayed instead of keeping a closure over the old <img>.
|
||||
const legacy = makeCard("asset-3", "c.png");
|
||||
const overlay = document.createElement("div");
|
||||
overlay.className = "asset-card-overlay";
|
||||
legacy.card.appendChild(overlay);
|
||||
|
||||
injectHoverCopyAction(overlay);
|
||||
const oldButton = overlay.querySelector(".saturnnodes-hover-copy");
|
||||
assert.ok(oldButton, "legacy overlay must receive a copy button");
|
||||
assert.equal(oldButton.dataset.saturnnodesAssetKey, "asset:asset-3", "button is bound to the current asset");
|
||||
|
||||
await oldButton.onclick({ preventDefault() {}, stopPropagation() {} });
|
||||
assert.deepEqual(copied, ["c.png"], "legacy button must copy the displayed image");
|
||||
|
||||
// Recycle the same overlay for another asset. The old button must be replaced,
|
||||
// not left bound to c.png.
|
||||
legacy.card.setAttribute("data-asset-id", "asset-4");
|
||||
legacy.img.src = "d.png";
|
||||
injectHoverCopyAction(overlay);
|
||||
const newButton = overlay.querySelector(".saturnnodes-hover-copy");
|
||||
assert.ok(newButton, "recycled overlay must still have a copy button");
|
||||
assert.notEqual(newButton, oldButton, "recycled overlay must rebind the copy button");
|
||||
assert.equal(newButton.dataset.saturnnodesAssetKey, "asset:asset-4", "button is rebound to the new asset");
|
||||
|
||||
await newButton.onclick({ preventDefault() {}, stopPropagation() {} });
|
||||
assert.deepEqual(copied, ["c.png", "d.png"], "rebound button must copy the new image");
|
||||
|
||||
// 5. Re-injecting for the same asset must not churn the button.
|
||||
injectHoverCopyAction(overlay);
|
||||
assert.equal(overlay.querySelector(".saturnnodes-hover-copy"), newButton, "same asset must keep the existing button");
|
||||
|
||||
console.log("copy prompt DOM regression test passed");
|
||||
@@ -0,0 +1,212 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import {
|
||||
FakeMutationObserver,
|
||||
createDocument,
|
||||
createWindow,
|
||||
stripImports,
|
||||
} from "./fake_dom.mjs";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "js", "copy_prompt.js");
|
||||
|
||||
const document = createDocument();
|
||||
const window = createWindow(document);
|
||||
|
||||
const extension = { name: "", beforeRegisterNodeDef: null };
|
||||
const app = {
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.beforeRegisterNodeDef = candidate.beforeRegisterNodeDef;
|
||||
},
|
||||
extensionManager: {
|
||||
setting: { get: () => true },
|
||||
toast: { add() {} },
|
||||
},
|
||||
ui: { settings: { get: () => undefined, getSettingValue: () => undefined } },
|
||||
graph: { setDirtyCanvas() {} },
|
||||
};
|
||||
|
||||
let source = stripImports(fs.readFileSync(sourcePath, "utf8"));
|
||||
source += `
|
||||
globalThis.__copyTest = {
|
||||
injectCopyPromptNextToDownload,
|
||||
injectContextMenuCopy,
|
||||
injectHoverCopyAction,
|
||||
registerAssetInteraction,
|
||||
getActiveImageSrc,
|
||||
scanAndInject,
|
||||
setCopyImagePrompt: (fn) => { copyImagePrompt = fn; },
|
||||
};
|
||||
`;
|
||||
|
||||
const copied = [];
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
document,
|
||||
window,
|
||||
navigator: {},
|
||||
app,
|
||||
api: { addEventListener() {}, fetchApi: async () => null },
|
||||
MutationObserver: FakeMutationObserver,
|
||||
setTimeout: () => 0,
|
||||
clearTimeout: () => {},
|
||||
setInterval: () => 0,
|
||||
clearInterval: () => {},
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
Date,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
Node: { ELEMENT_NODE: 1 },
|
||||
MouseEvent: class MouseEvent {
|
||||
constructor(type, init = {}) {
|
||||
this.type = type;
|
||||
this.bubbles = !!init.bubbles;
|
||||
this.cancelable = !!init.cancelable;
|
||||
}
|
||||
},
|
||||
});
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
|
||||
assert.equal(typeof extension.beforeRegisterNodeDef, "function", "copy prompt extension must register");
|
||||
|
||||
const copyTest = context.__copyTest;
|
||||
copyTest.setCopyImagePrompt(async src => {
|
||||
copied.push(src);
|
||||
return true;
|
||||
});
|
||||
|
||||
const observer = FakeMutationObserver.instances.find(
|
||||
item => item.target === document.body && item.options?.childList
|
||||
);
|
||||
assert.ok(observer, "copy prompt extension must observe restored DOM insertions");
|
||||
|
||||
function makeCard(assetId, src, { queue = false, icon = false, attach = true } = {}) {
|
||||
const card = document.createElement("div");
|
||||
card.className = "asset-card";
|
||||
if (assetId) card.setAttribute("data-asset-id", assetId);
|
||||
if (queue) card.setAttribute("data-job-id", "job-1");
|
||||
const img = document.createElement("img");
|
||||
img.src = src;
|
||||
if (icon) img.classList.add("size-8");
|
||||
card.appendChild(img);
|
||||
|
||||
const bar = document.createElement("div");
|
||||
const download = document.createElement("button");
|
||||
download.setAttribute("aria-label", "Download");
|
||||
bar.appendChild(download);
|
||||
card.appendChild(bar);
|
||||
|
||||
if (attach) document.body.appendChild(card);
|
||||
return { card, img, bar, download };
|
||||
}
|
||||
|
||||
function insertRestored(card) {
|
||||
observer.emit([{ type: "childList", target: document.body, addedNodes: [card], removedNodes: [] }]);
|
||||
}
|
||||
|
||||
// 1. A restored asset card inserted after page load receives a copy button
|
||||
// through the MutationObserver path (not only through the periodic scan).
|
||||
const restored = makeCard("restored-1", "restored.png", { attach: false });
|
||||
insertRestored(restored.card);
|
||||
const restoredButton = restored.bar.querySelector(".saturnnodes-hover-copy");
|
||||
assert.ok(restoredButton, "restored card must receive a copy button");
|
||||
await restoredButton.onclick({ preventDefault() {}, stopPropagation() {} });
|
||||
assert.deepEqual(copied, ["restored.png"], "restored copy button must copy the restored image");
|
||||
|
||||
// 2. Restored queue rows must never receive asset-card buttons.
|
||||
const queueRow = makeCard("queue-1", "queued.png", { queue: true, attach: false });
|
||||
insertRestored(queueRow.card);
|
||||
assert.equal(
|
||||
queueRow.bar.querySelector(".saturnnodes-hover-copy"),
|
||||
null,
|
||||
"restored queue rows must not receive asset copy buttons"
|
||||
);
|
||||
|
||||
// 3. Tiny icon images (download/status icons) must not be treated as assets.
|
||||
const iconCard = makeCard("icon-1", "icon.png", { icon: true, attach: false });
|
||||
insertRestored(iconCard.card);
|
||||
assert.equal(
|
||||
iconCard.bar.querySelector(".saturnnodes-hover-copy"),
|
||||
null,
|
||||
"icon-sized images must not receive asset copy buttons"
|
||||
);
|
||||
|
||||
// 4. A restored card reused for another asset must copy the asset currently
|
||||
// displayed, not the one captured when the button was injected.
|
||||
const recycled = makeCard("recycled-1", "first.png", { attach: false });
|
||||
insertRestored(recycled.card);
|
||||
const recycledButton = recycled.bar.querySelector(".saturnnodes-hover-copy");
|
||||
assert.ok(recycledButton, "recycled card must start with a copy button");
|
||||
|
||||
recycled.card.setAttribute("data-asset-id", "recycled-2");
|
||||
recycled.img.src = "second.png";
|
||||
await recycledButton.onclick({ preventDefault() {}, stopPropagation() {} });
|
||||
assert.deepEqual(copied, ["restored.png", "second.png"], "reused button must copy the live image");
|
||||
|
||||
// 5. Re-injecting for the same restored card must not duplicate the button.
|
||||
copyTest.injectCopyPromptNextToDownload(recycled.download);
|
||||
assert.equal(
|
||||
recycled.bar.querySelectorAll(".saturnnodes-hover-copy").length,
|
||||
1,
|
||||
"restored card must keep exactly one copy button"
|
||||
);
|
||||
|
||||
// 6. Context menus attached to restored cards inject one copy item and copy
|
||||
// the currently interacted asset.
|
||||
const menuCard = makeCard("menu-1", "menu-image.png");
|
||||
copyTest.registerAssetInteraction(menuCard.img);
|
||||
const menu = document.createElement("div");
|
||||
menu.className = "p-contextmenu";
|
||||
const downloadItem = document.createElement("li");
|
||||
downloadItem.setAttribute("role", "menuitem");
|
||||
downloadItem.setAttribute("aria-label", "Download");
|
||||
const downloadButton = document.createElement("button");
|
||||
downloadItem.appendChild(downloadButton);
|
||||
menu.appendChild(downloadItem);
|
||||
document.body.appendChild(menu);
|
||||
|
||||
copyTest.injectContextMenuCopy(menu);
|
||||
const copyItem = menu.querySelector(".saturnnodes-contextmenu-copy");
|
||||
assert.ok(copyItem, "restored context menu must receive a copy item");
|
||||
assert.equal(
|
||||
menu.querySelectorAll(".saturnnodes-contextmenu-copy").length,
|
||||
1,
|
||||
"context menu injection must be idempotent"
|
||||
);
|
||||
|
||||
const beforeMenuClick = copied.length;
|
||||
copyItem.dispatchEvent({ type: "click", preventDefault() {}, stopPropagation() {} });
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
assert.deepEqual(copied.slice(beforeMenuClick), ["menu-image.png"], "menu must copy the active asset");
|
||||
|
||||
// 7. A restored card without an asset id but with a stable image URL still
|
||||
// resolves its image for menu actions.
|
||||
const srcCard = makeCard(null, "src-only.png");
|
||||
copyTest.registerAssetInteraction(srcCard.img);
|
||||
assert.equal(copyTest.getActiveImageSrc(), "src-only.png");
|
||||
|
||||
// 8. The periodic scanner must inject into restored cards even when no
|
||||
// mutation was observed for them (e.g. restored before the observer ran).
|
||||
const scanned = makeCard("scanned-1", "scanned.png");
|
||||
copyTest.scanAndInject();
|
||||
assert.ok(
|
||||
scanned.bar.querySelector(".saturnnodes-hover-copy"),
|
||||
"periodic scan must cover restored cards"
|
||||
);
|
||||
|
||||
console.log("copy prompt restore DOM regression test passed");
|
||||
@@ -0,0 +1,38 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import threading
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.decision_node import DecisionManager, SaturnDecision
|
||||
|
||||
class TestDecisionNode(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = SaturnDecision()
|
||||
|
||||
def test_decision_manager_register_and_trigger(self):
|
||||
uid = "test_decision_123"
|
||||
event = threading.Event()
|
||||
DecisionManager.register_wait(uid, event)
|
||||
self.assertIn(uid, DecisionManager._waiting_events)
|
||||
|
||||
triggered = DecisionManager.trigger_action(uid, "cancel")
|
||||
self.assertTrue(triggered)
|
||||
self.assertTrue(event.is_set())
|
||||
|
||||
action = DecisionManager.get_action(uid)
|
||||
self.assertEqual(action, "cancel")
|
||||
self.assertNotIn(uid, DecisionManager._waiting_events)
|
||||
|
||||
def test_decision_disabled_mode(self):
|
||||
res = self.node.decide(disable=True, send_os_notification=False, timeout=-1, unique_id="disabled_node")
|
||||
self.assertEqual(res, (False,))
|
||||
|
||||
def test_decision_manager_unknown_id_default(self):
|
||||
action = DecisionManager.get_action("non_existent_id")
|
||||
self.assertEqual(action, "continue")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,133 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import {
|
||||
FakeMutationObserver,
|
||||
createDocument,
|
||||
createWindow,
|
||||
stripImports,
|
||||
} from "./fake_dom.mjs";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "floating_latent_dock.js");
|
||||
|
||||
const document = createDocument();
|
||||
const window = createWindow(document);
|
||||
|
||||
const extension = { name: "", setup: null };
|
||||
const app = {
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.setup = candidate.setup;
|
||||
},
|
||||
ui: { settings: { get: () => undefined, getSettingValue: () => undefined } },
|
||||
graph: { setDirtyCanvas() {} },
|
||||
};
|
||||
|
||||
let source = stripImports(fs.readFileSync(sourcePath, "utf8"));
|
||||
source += "\nglobalThis.__dockTest = { FloatingLatentDock };\n";
|
||||
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
document,
|
||||
window,
|
||||
app,
|
||||
api: { addEventListener() {} },
|
||||
localStorage: { getItem: () => null, setItem() {} },
|
||||
MutationObserver: FakeMutationObserver,
|
||||
setTimeout,
|
||||
clearTimeout,
|
||||
setInterval: () => 0,
|
||||
clearInterval: () => {},
|
||||
requestAnimationFrame: window.requestAnimationFrame,
|
||||
cancelAnimationFrame: window.cancelAnimationFrame,
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
Date,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
});
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
|
||||
assert.equal(typeof extension.setup, "function", "floating dock extension must register");
|
||||
const { FloatingLatentDock } = context.__dockTest;
|
||||
|
||||
const dock = {
|
||||
dockEl: document.createElement("div"),
|
||||
btnEl: null,
|
||||
updates: 0,
|
||||
requestPositionUpdate() {
|
||||
this.updates++;
|
||||
},
|
||||
};
|
||||
document.body.appendChild(dock.dockEl);
|
||||
|
||||
// 1. No minimap and no toolbar: dock sits at the fallback offset.
|
||||
FloatingLatentDock.prototype.updatePosition.call(dock);
|
||||
assert.equal(dock.dockEl.style.bottom, "54px", "fallback dock offset");
|
||||
|
||||
// 2. A minimap mounted deep in the tree moves the dock above it.
|
||||
const minimap = document.createElement("div");
|
||||
minimap.className = "minimap-main-container";
|
||||
minimap.setRect({ top: 700, height: 150, bottom: 850 });
|
||||
minimap.offsetParent = {};
|
||||
document.body.appendChild(minimap);
|
||||
|
||||
FloatingLatentDock.prototype.updatePosition.call(dock);
|
||||
assert.equal(dock.dockEl.style.bottom, "208px", "dock must sit 8px above the minimap");
|
||||
|
||||
// 3. The observer must react to a deep minimap insertion/removal, not only to
|
||||
// direct children of <body>.
|
||||
FloatingLatentDock.prototype.setupPositionObserver.call(dock);
|
||||
const observer = FakeMutationObserver.instances.find(
|
||||
item => item.options?.childList && item.options?.subtree
|
||||
);
|
||||
assert.ok(observer, "dock observer must watch the whole subtree");
|
||||
|
||||
observer.emit([{
|
||||
type: "childList",
|
||||
target: document.body,
|
||||
addedNodes: [minimap],
|
||||
removedNodes: [],
|
||||
}]);
|
||||
assert.equal(dock.updates, 1, "minimap insertion must request a position update");
|
||||
|
||||
observer.emit([{
|
||||
type: "childList",
|
||||
target: document.body,
|
||||
addedNodes: [document.createElement("div")],
|
||||
removedNodes: [],
|
||||
}]);
|
||||
assert.equal(dock.updates, 1, "unrelated DOM insertions must not request updates");
|
||||
|
||||
observer.emit([{
|
||||
type: "childList",
|
||||
target: document.body,
|
||||
addedNodes: [],
|
||||
removedNodes: [minimap],
|
||||
}]);
|
||||
assert.equal(dock.updates, 2, "minimap removal must request a position update");
|
||||
|
||||
// 4. Removing the minimap restores the toolbar-based offset.
|
||||
minimap.remove();
|
||||
const toolbar = document.createElement("div");
|
||||
toolbar.setAttribute("role", "toolbar");
|
||||
toolbar.setAttribute("aria-label", "Canvas Toolbar");
|
||||
toolbar.setRect({ top: 820, height: 40, bottom: 860 });
|
||||
document.body.appendChild(toolbar);
|
||||
|
||||
FloatingLatentDock.prototype.updatePosition.call(dock);
|
||||
assert.equal(dock.dockEl.style.bottom, "88px", "dock must sit 8px above the canvas toolbar");
|
||||
|
||||
console.log("floating dock DOM regression test passed");
|
||||
@@ -0,0 +1,46 @@
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import unittest
|
||||
|
||||
|
||||
DOM_TESTS = [
|
||||
"test_power_popup_dom.mjs",
|
||||
"test_copy_prompt_dom.mjs",
|
||||
"test_copy_prompt_restore_dom.mjs",
|
||||
"test_floating_dock_dom.mjs",
|
||||
"test_widget_guards_dom.mjs",
|
||||
"test_queue_height_dom.mjs",
|
||||
"test_batch_queue_dom.mjs",
|
||||
"test_batch_queue_snapshot_dom.mjs",
|
||||
"test_batch_queue_restore_dom.mjs",
|
||||
"test_persistent_queue_claim_dom.mjs",
|
||||
]
|
||||
|
||||
|
||||
class TestFrontendDomRegressions(unittest.TestCase):
|
||||
def test_dom_regressions(self):
|
||||
node = shutil.which("node")
|
||||
if not node:
|
||||
self.skipTest("Node.js is not available")
|
||||
|
||||
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
for script_name in DOM_TESTS:
|
||||
with self.subTest(script=script_name):
|
||||
script = os.path.join(repo_root, "tests", script_name)
|
||||
result = subprocess.run(
|
||||
[node, script],
|
||||
cwd=repo_root,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=60,
|
||||
)
|
||||
self.assertEqual(
|
||||
result.returncode,
|
||||
0,
|
||||
msg=f"{script_name} failed:\n{result.stdout}\n{result.stderr}",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,46 @@
|
||||
import os
|
||||
import re
|
||||
import unittest
|
||||
|
||||
|
||||
WEB_ROOT = os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "web"
|
||||
)
|
||||
|
||||
# ComfyUI's settings API warns every time getSettingValue() receives a default
|
||||
# value as its second argument. Settings definitions already carry defaults, so
|
||||
# the argument is redundant and floods the browser console when called per item.
|
||||
DEPRECATED_SECOND_ARGUMENT = re.compile(
|
||||
r"""getSettingValue\s*\??\.?\s*\(\s*(?:"[^"]*"|'[^']*'|`[^`]*`)\s*,"""
|
||||
)
|
||||
|
||||
|
||||
class TestFrontendSettingsApi(unittest.TestCase):
|
||||
def test_no_deprecated_default_value_argument(self):
|
||||
offenders = []
|
||||
for root, _dirs, files in os.walk(WEB_ROOT):
|
||||
for name in sorted(files):
|
||||
if not name.endswith(".js"):
|
||||
continue
|
||||
path = os.path.join(root, name)
|
||||
with open(path, encoding="utf-8") as handle:
|
||||
source = handle.read()
|
||||
for match in DEPRECATED_SECOND_ARGUMENT.finditer(source):
|
||||
line_number = source.count("\n", 0, match.start()) + 1
|
||||
relative = os.path.relpath(path, WEB_ROOT)
|
||||
offenders.append(f"{relative}:{line_number}")
|
||||
|
||||
self.assertEqual(
|
||||
offenders,
|
||||
[],
|
||||
msg=(
|
||||
"getSettingValue(id, defaultValue) is deprecated in ComfyUI and "
|
||||
"logs a console warning on every call. Read the setting default "
|
||||
"from its definition instead: "
|
||||
+ ", ".join(offenders)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,239 @@
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper # noqa: F401 (redirects SaturnNodes state to a temp dir)
|
||||
|
||||
import nodes.generation_times as generation_times
|
||||
|
||||
|
||||
class TestGenerationTimesExtended(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.old_file = generation_times.GENERATION_TIMES_FILE
|
||||
self.old_user_dir = generation_times.USER_DIR
|
||||
self.old_data = generation_times._DATA
|
||||
self.old_mtime = generation_times._DATA_MTIME
|
||||
self.old_patched = generation_times._PATCHED
|
||||
self.old_starts = dict(generation_times._START_TIMES)
|
||||
generation_times.GENERATION_TIMES_FILE = os.path.join(self.temp_dir, "generation_times.json")
|
||||
generation_times.USER_DIR = self.temp_dir
|
||||
generation_times._DATA = None
|
||||
generation_times._DATA_MTIME = None
|
||||
generation_times._PATCHED = False
|
||||
generation_times._START_TIMES.clear()
|
||||
|
||||
def tearDown(self):
|
||||
generation_times.GENERATION_TIMES_FILE = self.old_file
|
||||
generation_times.USER_DIR = self.old_user_dir
|
||||
generation_times._DATA = self.old_data
|
||||
generation_times._DATA_MTIME = self.old_mtime
|
||||
generation_times._PATCHED = self.old_patched
|
||||
generation_times._START_TIMES.clear()
|
||||
generation_times._START_TIMES.update(self.old_starts)
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def _write_payload(self, payload):
|
||||
with open(generation_times.GENERATION_TIMES_FILE, "w", encoding="utf-8") as handle:
|
||||
json.dump(payload, handle)
|
||||
generation_times._DATA = None
|
||||
generation_times._DATA_MTIME = None
|
||||
|
||||
def test_rejects_non_finite_negative_and_non_numeric_durations(self):
|
||||
bad_values = [float("nan"), float("inf"), float("-inf"), -0.001, "not-a-number", None, object()]
|
||||
for index, value in enumerate(bad_values):
|
||||
generation_times.record_generation_time(
|
||||
f"pid-{index}",
|
||||
value,
|
||||
[{"filename": f"bad{index}.png", "subfolder": ""}],
|
||||
)
|
||||
|
||||
self.assertFalse(
|
||||
os.path.exists(generation_times.GENERATION_TIMES_FILE),
|
||||
"invalid durations must not create a generation-times file",
|
||||
)
|
||||
for index in range(len(bad_values)):
|
||||
self.assertIsNone(generation_times.get_generation_time(f"bad{index}.png"))
|
||||
|
||||
def test_rejects_missing_or_malformed_file_entries(self):
|
||||
generation_times.record_generation_time("pid", 2.0, None)
|
||||
generation_times.record_generation_time("pid", 2.0, [])
|
||||
generation_times.record_generation_time("pid", 2.0, [None, "x", {}, {"filename": ""}])
|
||||
self.assertIsNone(generation_times.get_generation_time(""))
|
||||
|
||||
generation_times.record_generation_time("pid", 2.0, [{"filename": "valid.png"}])
|
||||
self.assertEqual(generation_times.get_generation_time("valid.png"), 2.0)
|
||||
|
||||
def test_rounds_and_overwrites_duplicate_file_keys(self):
|
||||
generation_times.record_generation_time(
|
||||
"pid-1", 1.23456, [{"filename": "same.png", "subfolder": "sub"}]
|
||||
)
|
||||
self.assertEqual(generation_times.get_generation_time("same.png", "sub"), 1.235)
|
||||
|
||||
generation_times.record_generation_time(
|
||||
"pid-2", 9.0, [{"filename": "same.png", "subfolder": "sub"}]
|
||||
)
|
||||
with open(generation_times.GENERATION_TIMES_FILE, "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
self.assertEqual(len(payload["entries"]), 1)
|
||||
self.assertEqual(payload["entries"]["sub/same.png"]["duration"], 9.0)
|
||||
self.assertEqual(payload["entries"]["sub/same.png"]["prompt_id"], "pid-2")
|
||||
|
||||
def test_pruning_keeps_the_newest_entries(self):
|
||||
times = iter([1000.0, 1001.0, 1002.0, 1003.0, 1004.0])
|
||||
with patch.object(generation_times, "MAX_ENTRIES", 3), \
|
||||
patch("nodes.generation_times.time.time", side_effect=lambda: next(times)):
|
||||
for index in range(5):
|
||||
generation_times.record_generation_time(
|
||||
f"pid-{index}",
|
||||
float(index),
|
||||
[{"filename": f"file{index}.png", "subfolder": ""}],
|
||||
)
|
||||
|
||||
with open(generation_times.GENERATION_TIMES_FILE, "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
self.assertEqual(
|
||||
set(payload["entries"]),
|
||||
{"file2.png", "file3.png", "file4.png"},
|
||||
)
|
||||
|
||||
def test_entry_timestamp_handles_malformed_values(self):
|
||||
self.assertEqual(generation_times._entry_timestamp(None), 0)
|
||||
self.assertEqual(generation_times._entry_timestamp("x"), 0)
|
||||
self.assertEqual(generation_times._entry_timestamp({}), 0)
|
||||
self.assertEqual(generation_times._entry_timestamp({"completed": "bad"}), 0)
|
||||
self.assertEqual(generation_times._entry_timestamp({"completed": -10}), 0)
|
||||
self.assertEqual(generation_times._entry_timestamp({"completed": float("inf")}), 0)
|
||||
self.assertEqual(generation_times._entry_timestamp({"completed": "1234"}), 1234)
|
||||
|
||||
def test_get_generation_time_rejects_corrupt_stored_durations(self):
|
||||
self._write_payload({
|
||||
"version": 1,
|
||||
"entries": {
|
||||
"nan.png": {"duration": float("nan"), "completed": 1},
|
||||
"inf.png": {"duration": float("inf"), "completed": 1},
|
||||
"neg.png": {"duration": -4, "completed": 1},
|
||||
"text.png": {"duration": "fast", "completed": 1},
|
||||
"ok.png": {"duration": "1.5", "completed": 1},
|
||||
},
|
||||
})
|
||||
for name in ("nan.png", "inf.png", "neg.png", "text.png"):
|
||||
self.assertIsNone(generation_times.get_generation_time(name), name)
|
||||
self.assertEqual(generation_times.get_generation_time("ok.png"), 1.5)
|
||||
|
||||
def test_load_data_recovers_from_corrupt_or_legacy_payloads(self):
|
||||
with open(generation_times.GENERATION_TIMES_FILE, "w", encoding="utf-8") as handle:
|
||||
handle.write("{ this is not json")
|
||||
data = generation_times._load_data()
|
||||
self.assertEqual(data, {"version": 1, "entries": {}})
|
||||
|
||||
self._write_payload({"entries": {"legacy.png": {"duration": 2.0, "completed": 5}}})
|
||||
data = generation_times._load_data()
|
||||
self.assertEqual(data["version"], 1)
|
||||
self.assertEqual(generation_times.get_generation_time("legacy.png"), 2.0)
|
||||
|
||||
def test_corrupt_entries_are_dropped_on_load_and_save(self):
|
||||
self._write_payload({
|
||||
"version": 1,
|
||||
"entries": {
|
||||
"nan.png": {"duration": float("nan"), "completed": 1},
|
||||
"inf.png": {"duration": float("inf"), "completed": 1},
|
||||
"neg.png": {"duration": -4, "completed": 1},
|
||||
"text.png": {"duration": "fast", "completed": 1},
|
||||
"ok.png": {"duration": "1.5", "completed": 1},
|
||||
"nan-completed.png": {"duration": 2.0, "completed": float("nan")},
|
||||
},
|
||||
})
|
||||
|
||||
data = generation_times._load_data()
|
||||
self.assertEqual(set(data["entries"]), {"ok.png", "nan-completed.png"})
|
||||
self.assertEqual(data["entries"]["ok.png"]["duration"], 1.5)
|
||||
self.assertEqual(data["entries"]["nan-completed.png"]["completed"], 0)
|
||||
|
||||
# A corrupt entry must not make every future save fail under
|
||||
# allow_nan=False. The valid entry and the new one must persist.
|
||||
generation_times.record_generation_time(
|
||||
"pid", 3.0, [{"filename": "new.png", "subfolder": ""}]
|
||||
)
|
||||
with open(generation_times.GENERATION_TIMES_FILE, "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
self.assertEqual(
|
||||
set(payload["entries"]),
|
||||
{"ok.png", "nan-completed.png", "new.png"},
|
||||
)
|
||||
self.assertEqual(payload["entries"]["new.png"]["duration"], 3.0)
|
||||
self.assertEqual(payload["entries"]["nan-completed.png"]["completed"], 0)
|
||||
|
||||
def test_save_retries_windows_replace_lock_and_cleans_temp_file(self):
|
||||
real_replace = os.replace
|
||||
attempts = {"count": 0}
|
||||
|
||||
def flaky_replace(src, dst):
|
||||
attempts["count"] += 1
|
||||
if attempts["count"] < 3:
|
||||
raise PermissionError("file is locked")
|
||||
return real_replace(src, dst)
|
||||
|
||||
with patch("nodes.generation_times.os.replace", side_effect=flaky_replace), \
|
||||
patch("nodes.generation_times.time.sleep"):
|
||||
generation_times._save_data({"version": 1, "entries": {"a.png": {"duration": 1}}})
|
||||
|
||||
self.assertEqual(attempts["count"], 3)
|
||||
self.assertTrue(os.path.exists(generation_times.GENERATION_TIMES_FILE))
|
||||
self.assertEqual(
|
||||
[name for name in os.listdir(self.temp_dir) if name.endswith(".tmp")],
|
||||
[],
|
||||
)
|
||||
|
||||
def test_save_failure_does_not_leave_temp_file_behind(self):
|
||||
with patch("nodes.generation_times.os.replace", side_effect=PermissionError("locked")), \
|
||||
patch("nodes.generation_times.time.sleep"):
|
||||
generation_times._save_data({"version": 1, "entries": {}})
|
||||
|
||||
self.assertEqual(
|
||||
[name for name in os.listdir(self.temp_dir) if name.endswith(".tmp")],
|
||||
[],
|
||||
)
|
||||
|
||||
def test_moved_file_fallback_handles_invalid_mtime(self):
|
||||
now = 1_700_000_000.0
|
||||
with patch("nodes.generation_times.time.time", return_value=now):
|
||||
generation_times.record_generation_time(
|
||||
"pid", 6.0, [{"filename": "moved.png", "subfolder": "old"}]
|
||||
)
|
||||
# A non-numeric mtime must not crash the fallback path.
|
||||
self.assertEqual(
|
||||
generation_times.get_generation_time("moved.png", "new", "not-a-number"),
|
||||
6.0,
|
||||
)
|
||||
|
||||
def test_extract_output_files_ignores_malformed_history(self):
|
||||
self.assertEqual(generation_times._extract_output_files(None), [])
|
||||
self.assertEqual(generation_times._extract_output_files({}), [])
|
||||
self.assertEqual(
|
||||
generation_times._extract_output_files({"outputs": {"1": {"images": "bad"}}}),
|
||||
[],
|
||||
)
|
||||
files = generation_times._extract_output_files({
|
||||
"outputs": {
|
||||
"1": {"images": [
|
||||
None,
|
||||
{"filename": ""},
|
||||
{"filename": "temp.png", "type": "temp"},
|
||||
{"filename": "saved.png", "subfolder": "sub"},
|
||||
]},
|
||||
"2": "bad",
|
||||
}
|
||||
})
|
||||
self.assertEqual(files, [{"filename": "saved.png", "subfolder": "sub", "type": "output"}])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,98 @@
|
||||
import atexit
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
SATURNNODES_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
# Keep every module-level state file (persistent queue, LoRA usage, caches,
|
||||
# prompt iterator state, generation times) inside a disposable test directory.
|
||||
# This is critical when ComfyUI is running concurrently: tests must never write
|
||||
# to the live ComfyUI/user/default/SaturnNodes directory.
|
||||
TEST_USER_ROOT = tempfile.mkdtemp(prefix="saturnnodes_test_user_")
|
||||
atexit.register(shutil.rmtree, TEST_USER_ROOT, ignore_errors=True)
|
||||
|
||||
# Ensure SaturnNodes root is prioritized at the very start of sys.path
|
||||
while SATURNNODES_ROOT in sys.path:
|
||||
sys.path.remove(SATURNNODES_ROOT)
|
||||
sys.path.insert(0, SATURNNODES_ROOT)
|
||||
|
||||
# If 'nodes' in sys.modules is ComfyUI's root nodes.py file, replace it with SaturnNodes' nodes package
|
||||
if "nodes" in sys.modules:
|
||||
nodes_mod = sys.modules["nodes"]
|
||||
if not hasattr(nodes_mod, "__path__") or not getattr(nodes_mod, "__file__", "").startswith(SATURNNODES_ROOT):
|
||||
del sys.modules["nodes"]
|
||||
|
||||
# Provide lightweight mock stubs for ComfyUI environment if running in standalone CI / test runner
|
||||
if "aiohttp" not in sys.modules:
|
||||
try:
|
||||
import aiohttp
|
||||
except ImportError:
|
||||
aiohttp_mock = types.ModuleType("aiohttp")
|
||||
web_mock = types.ModuleType("aiohttp.web")
|
||||
web_mock.json_response = MagicMock(return_value={})
|
||||
web_mock.Response = MagicMock()
|
||||
aiohttp_mock.web = web_mock
|
||||
sys.modules["aiohttp"] = aiohttp_mock
|
||||
sys.modules["aiohttp.web"] = web_mock
|
||||
|
||||
if "piexif" not in sys.modules:
|
||||
try:
|
||||
import piexif
|
||||
except ImportError:
|
||||
piexif_mock = types.ModuleType("piexif")
|
||||
piexif_mock.load = MagicMock(return_value={})
|
||||
piexif_mock.dump = MagicMock(return_value=b"")
|
||||
piexif_mock.insert = MagicMock()
|
||||
sys.modules["piexif"] = piexif_mock
|
||||
|
||||
if "pystray" not in sys.modules:
|
||||
try:
|
||||
import pystray
|
||||
except ImportError:
|
||||
pystray_mock = types.ModuleType("pystray")
|
||||
sys.modules["pystray"] = pystray_mock
|
||||
|
||||
if "comfy" not in sys.modules:
|
||||
comfy_mock = types.ModuleType("comfy")
|
||||
mm_mock = types.ModuleType("comfy.model_management")
|
||||
mm_mock.throw_exception_if_processing_interrupted = MagicMock()
|
||||
comfy_mock.model_management = mm_mock
|
||||
|
||||
sd_mock = types.ModuleType("comfy.sd")
|
||||
sd_mock.load_lora_for_models = MagicMock(side_effect=lambda m, c, l, sm, sc: (m, c))
|
||||
comfy_mock.sd = sd_mock
|
||||
|
||||
utils_mock = types.ModuleType("comfy.utils")
|
||||
utils_mock.load_torch_file = MagicMock(return_value={})
|
||||
comfy_mock.utils = utils_mock
|
||||
|
||||
sys.modules["comfy"] = comfy_mock
|
||||
sys.modules["comfy.model_management"] = mm_mock
|
||||
sys.modules["comfy.sd"] = sd_mock
|
||||
sys.modules["comfy.utils"] = utils_mock
|
||||
|
||||
if "folder_paths" not in sys.modules:
|
||||
fp_mock = types.ModuleType("folder_paths")
|
||||
fp_mock.get_input_directory = MagicMock(return_value=os.path.join(SATURNNODES_ROOT, "input"))
|
||||
fp_mock.get_output_directory = MagicMock(return_value=os.path.join(SATURNNODES_ROOT, "output"))
|
||||
fp_mock.get_temp_directory = MagicMock(return_value=os.path.join(SATURNNODES_ROOT, "temp"))
|
||||
fp_mock.get_filename_list = MagicMock(return_value=[])
|
||||
fp_mock.get_full_path = MagicMock(return_value="")
|
||||
fp_mock.base_path = SATURNNODES_ROOT
|
||||
sys.modules["folder_paths"] = fp_mock
|
||||
|
||||
# Make get_saturnnodes_user_dir() resolve to the disposable test directory,
|
||||
# whether folder_paths is the lightweight mock or the real ComfyUI module.
|
||||
import folder_paths
|
||||
folder_paths.get_user_directory = MagicMock(return_value=TEST_USER_ROOT)
|
||||
|
||||
if "server" not in sys.modules:
|
||||
server_mock = types.ModuleType("server")
|
||||
ps_mock = MagicMock()
|
||||
ps_mock.instance.routes = MagicMock()
|
||||
server_mock.PromptServer = ps_mock
|
||||
sys.modules["server"] = server_mock
|
||||
@@ -0,0 +1,245 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from nodes.image_loader import trace_positive_prompt_from_graph, parse_positive_from_parameters
|
||||
import nodes.image_loader as image_loader
|
||||
|
||||
class TestPromptExtraction(unittest.TestCase):
|
||||
def test_standard_ksampler_graph(self):
|
||||
graph = {
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": ["6", 0],
|
||||
"negative": ["7", 0],
|
||||
"model": ["4", 0]
|
||||
}
|
||||
},
|
||||
"6": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {
|
||||
"text": "a majestic lion in the savanna at golden hour"
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {
|
||||
"text": "ugly, blurry, low quality"
|
||||
}
|
||||
}
|
||||
}
|
||||
res = trace_positive_prompt_from_graph(graph)
|
||||
self.assertEqual(res, "a majestic lion in the savanna at golden hour")
|
||||
|
||||
def test_flux_guidance_graph(self):
|
||||
graph = {
|
||||
"1": {
|
||||
"class_type": "KSamplerSelect",
|
||||
"inputs": {}
|
||||
},
|
||||
"2": {
|
||||
"class_type": "SamplerCustomAdvanced",
|
||||
"inputs": {
|
||||
"guider": ["3", 0]
|
||||
}
|
||||
},
|
||||
"3": {
|
||||
"class_type": "BasicGuider",
|
||||
"inputs": {
|
||||
"conditioning": ["4", 0]
|
||||
}
|
||||
},
|
||||
"4": {
|
||||
"class_type": "FluxGuidance",
|
||||
"inputs": {
|
||||
"guidance": 3.5,
|
||||
"conditioning": ["5", 0]
|
||||
}
|
||||
},
|
||||
"5": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {
|
||||
"text": "a cyberpunk futuristic street in Tokyo, neon reflections"
|
||||
}
|
||||
}
|
||||
}
|
||||
res = trace_positive_prompt_from_graph(graph)
|
||||
self.assertEqual(res, "a cyberpunk futuristic street in Tokyo, neon reflections")
|
||||
|
||||
def test_sdxl_text_g_and_text_l(self):
|
||||
graph = {
|
||||
"10": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": ["20", 0],
|
||||
"negative": ["21", 0]
|
||||
}
|
||||
},
|
||||
"20": {
|
||||
"class_type": "CLIPTextEncodeSDXL",
|
||||
"inputs": {
|
||||
"text_g": "an astronaut riding a horse",
|
||||
"text_l": "photorealistic, 8k resolution"
|
||||
}
|
||||
},
|
||||
"21": {
|
||||
"class_type": "CLIPTextEncodeSDXL",
|
||||
"inputs": {
|
||||
"text_g": "blurry",
|
||||
"text_l": "cartoon"
|
||||
}
|
||||
}
|
||||
}
|
||||
res = trace_positive_prompt_from_graph(graph)
|
||||
self.assertIn("an astronaut riding a horse", res)
|
||||
self.assertIn("photorealistic, 8k resolution", res)
|
||||
self.assertNotIn("blurry", res)
|
||||
|
||||
def test_conditioning_combine(self):
|
||||
graph = {
|
||||
"1": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": ["2", 0],
|
||||
"negative": ["9", 0]
|
||||
}
|
||||
},
|
||||
"2": {
|
||||
"class_type": "ConditioningCombine",
|
||||
"inputs": {
|
||||
"conditioning_1": ["3", 0],
|
||||
"conditioning_2": ["4", 0]
|
||||
}
|
||||
},
|
||||
"3": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "part 1 of positive prompt"}
|
||||
},
|
||||
"4": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "part 2 of positive prompt"}
|
||||
}
|
||||
}
|
||||
res = trace_positive_prompt_from_graph(graph)
|
||||
self.assertIn("part 1 of positive prompt", res)
|
||||
self.assertIn("part 2 of positive prompt", res)
|
||||
|
||||
def test_primitive_string_link(self):
|
||||
graph = {
|
||||
"1": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": ["2", 0]
|
||||
}
|
||||
},
|
||||
"2": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {
|
||||
"text": ["3", 0]
|
||||
}
|
||||
},
|
||||
"3": {
|
||||
"class_type": "PrimitiveString",
|
||||
"inputs": {
|
||||
"value": "string from primitive node"
|
||||
}
|
||||
}
|
||||
}
|
||||
res = trace_positive_prompt_from_graph(graph)
|
||||
self.assertEqual(res, "string from primitive node")
|
||||
|
||||
def test_a1111_parameters_parsing(self):
|
||||
param_str = """A portrait of an ancient wizard with a glowing crystal staff
|
||||
Negative prompt: ugly, deformed, bad anatomy
|
||||
Steps: 30, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 12345, Size: 512x768"""
|
||||
res = parse_positive_from_parameters(param_str)
|
||||
self.assertEqual(res, "A portrait of an ancient wizard with a glowing crystal staff")
|
||||
|
||||
class TestImageFolderScanning(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.folder = os.path.join(self.temp_dir, "images")
|
||||
os.makedirs(self.folder)
|
||||
self.old_cache = image_loader.CACHE_FILE
|
||||
self.old_thumbs = image_loader.IMAGE_THUMBNAILS_DIR
|
||||
image_loader.CACHE_FILE = os.path.join(self.temp_dir, "prompts_cache.json")
|
||||
image_loader.IMAGE_THUMBNAILS_DIR = os.path.join(self.temp_dir, "thumbs")
|
||||
|
||||
def tearDown(self):
|
||||
image_loader.CACHE_FILE = self.old_cache
|
||||
image_loader.IMAGE_THUMBNAILS_DIR = self.old_thumbs
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def _make_image(self, name, size=(32, 32), mtime=None):
|
||||
path = os.path.join(self.folder, name)
|
||||
Image.new("RGB", size, (120, 90, 60)).save(path)
|
||||
if mtime:
|
||||
os.utime(path, (mtime, mtime))
|
||||
return path
|
||||
|
||||
def test_metadata_limit_and_cache_reuse(self):
|
||||
base = 1_600_000_000
|
||||
self._make_image("a.png", mtime=base)
|
||||
self._make_image("b.png", mtime=base + 10)
|
||||
self._make_image("c.png", mtime=base + 20)
|
||||
|
||||
calls = []
|
||||
|
||||
def fake_extract(path):
|
||||
calls.append(path)
|
||||
return (f"prompt-{os.path.basename(path)}", 32, 32)
|
||||
|
||||
with patch("nodes.image_loader.extract_metadata_from_image", side_effect=fake_extract):
|
||||
names, mapping, prompts = image_loader.scan_image_folder(self.folder, metadata_limit=2)
|
||||
self.assertEqual(names, ["c.png", "b.png", "a.png"])
|
||||
self.assertEqual(mapping["a.png"], "a.png")
|
||||
self.assertEqual(len(calls), 2)
|
||||
self.assertEqual(prompts["c.png"], "prompt-c.png")
|
||||
self.assertEqual(prompts["b.png"], "prompt-b.png")
|
||||
self.assertEqual(prompts["a.png"], "")
|
||||
|
||||
# Unchanged mtimes must be served from the disk cache.
|
||||
calls.clear()
|
||||
image_loader.scan_image_folder(self.folder, metadata_limit=2)
|
||||
self.assertEqual(calls, [])
|
||||
|
||||
self.assertTrue(os.path.exists(image_loader.CACHE_FILE))
|
||||
|
||||
def test_thumbnail_disk_cache_and_mtime_invalidation(self):
|
||||
path = self._make_image("large.png", size=(640, 480), mtime=1_600_000_000)
|
||||
first = image_loader.get_or_create_image_thumbnail(path, 1_600_000_000, 256)
|
||||
self.assertTrue(first.endswith(".webp"))
|
||||
self.assertTrue(os.path.exists(first))
|
||||
with Image.open(first) as thumb:
|
||||
self.assertLessEqual(max(thumb.size), 256)
|
||||
|
||||
second = image_loader.get_or_create_image_thumbnail(path, 1_600_000_000, 256)
|
||||
self.assertEqual(first, second)
|
||||
|
||||
third = image_loader.get_or_create_image_thumbnail(path, 1_600_000_100, 256)
|
||||
self.assertNotEqual(first, third)
|
||||
self.assertTrue(os.path.exists(third))
|
||||
|
||||
def test_metadata_scan_limit_setting_is_bounded(self):
|
||||
with patch("nodes.image_loader.get_env_setting", return_value="42"):
|
||||
self.assertEqual(image_loader.get_image_metadata_scan_limit(), 42)
|
||||
with patch("nodes.image_loader.get_env_setting", return_value="oops"):
|
||||
self.assertEqual(image_loader.get_image_metadata_scan_limit(), 128)
|
||||
with patch("nodes.image_loader.get_env_setting", return_value="-5"):
|
||||
self.assertEqual(image_loader.get_image_metadata_scan_limit(), 0)
|
||||
with patch("nodes.image_loader.get_env_setting", return_value="99999"):
|
||||
self.assertEqual(image_loader.get_image_metadata_scan_limit(), 2000)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,200 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper # noqa: F401
|
||||
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
import nodes.image_loader as image_loader
|
||||
from nodes.image_loader import (
|
||||
decode_bytes_safely,
|
||||
parse_positive_from_parameters,
|
||||
scan_image_folder,
|
||||
trace_positive_prompt_from_graph,
|
||||
)
|
||||
|
||||
|
||||
class TestImageLoaderExtended(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.cache_file = os.path.join(self.temp_dir, "image_prompts_cache.json")
|
||||
self.old_cache = image_loader.CACHE_FILE
|
||||
image_loader.CACHE_FILE = self.cache_file
|
||||
|
||||
def tearDown(self):
|
||||
image_loader.CACHE_FILE = self.old_cache
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def _make_png(self, name, mtime, subfolder="", prompt_graph=None, parameters=None):
|
||||
folder = os.path.join(self.temp_dir, subfolder) if subfolder else self.temp_dir
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
path = os.path.join(folder, name)
|
||||
info = PngImagePlugin.PngInfo()
|
||||
if prompt_graph is not None:
|
||||
info.add_text("prompt", json.dumps(prompt_graph))
|
||||
if parameters is not None:
|
||||
info.add_text("parameters", parameters)
|
||||
Image.new("RGB", (6, 6), (30, 40, 50)).save(path, pnginfo=info)
|
||||
os.utime(path, (mtime, mtime))
|
||||
return path
|
||||
|
||||
def test_equal_mtimes_are_sorted_deterministically(self):
|
||||
base = 1_700_000_000
|
||||
self._make_png("b.png", base)
|
||||
self._make_png("a.png", base)
|
||||
self._make_png("c.png", base, subfolder="sub")
|
||||
self._make_png("a.png", base, subfolder="sub")
|
||||
|
||||
names, mapping, prompts = scan_image_folder(self.temp_dir, metadata_limit=0)
|
||||
expected = ["a.png", "b.png", "sub/a.png", "sub/c.png"]
|
||||
self.assertEqual(names, expected)
|
||||
self.assertEqual(list(mapping), expected)
|
||||
self.assertEqual(list(prompts), expected)
|
||||
self.assertTrue(all(value == "" for value in prompts.values()))
|
||||
|
||||
def test_newest_images_come_first_across_subfolders(self):
|
||||
base = 1_700_000_000
|
||||
self._make_png("old.png", base)
|
||||
self._make_png("new.png", base + 50, subfolder="nested")
|
||||
|
||||
names, _, _ = scan_image_folder(self.temp_dir, metadata_limit=0)
|
||||
self.assertEqual(names, ["nested/new.png", "old.png"])
|
||||
|
||||
def test_metadata_limit_zero_skips_extraction(self):
|
||||
self._make_png("one.png", 1_700_000_000)
|
||||
with patch("nodes.image_loader.extract_metadata_from_image") as extract:
|
||||
names, _, prompts = scan_image_folder(self.temp_dir, metadata_limit=0)
|
||||
extract.assert_not_called()
|
||||
self.assertEqual(names, ["one.png"])
|
||||
self.assertEqual(prompts, {"one.png": ""})
|
||||
|
||||
def test_metadata_limit_only_reads_newest_entries(self):
|
||||
base = 1_700_000_000
|
||||
self._make_png("older.png", base)
|
||||
self._make_png("newer.png", base + 10)
|
||||
|
||||
with patch(
|
||||
"nodes.image_loader.extract_metadata_from_image",
|
||||
side_effect=lambda path: ("prompt:" + os.path.basename(path), 6, 6),
|
||||
) as extract:
|
||||
names, _, prompts = scan_image_folder(self.temp_dir, metadata_limit=1)
|
||||
|
||||
self.assertEqual(names, ["newer.png", "older.png"])
|
||||
self.assertEqual(extract.call_count, 1)
|
||||
self.assertEqual(prompts["newer.png"], "prompt:newer.png")
|
||||
self.assertEqual(prompts["older.png"], "")
|
||||
|
||||
def test_metadata_cache_is_reused_and_invalidated_by_mtime(self):
|
||||
path = self._make_png("cached.png", 1_700_000_000)
|
||||
with patch(
|
||||
"nodes.image_loader.extract_metadata_from_image",
|
||||
return_value=("cached prompt", 6, 6),
|
||||
) as extract:
|
||||
scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
self.assertEqual(extract.call_count, 1)
|
||||
|
||||
os.utime(path, (1_700_000_100, 1_700_000_100))
|
||||
scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
self.assertEqual(extract.call_count, 2)
|
||||
|
||||
def test_png_prompt_metadata_is_extracted(self):
|
||||
graph = {
|
||||
"1": {"class_type": "KSampler", "inputs": {"positive": ["2", 0], "negative": ["3", 0]}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "a beautiful castle"}},
|
||||
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "ugly"}},
|
||||
}
|
||||
self._make_png("graph.png", 1_700_000_000, prompt_graph=graph)
|
||||
names, _, prompts = scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
self.assertEqual(names, ["graph.png"])
|
||||
self.assertEqual(prompts["graph.png"], "a beautiful castle")
|
||||
|
||||
def test_corrupt_image_does_not_break_the_scan(self):
|
||||
path = os.path.join(self.temp_dir, "broken.png")
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(b"not a png")
|
||||
os.utime(path, (1_700_000_000, 1_700_000_000))
|
||||
|
||||
names, _, prompts = scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
self.assertEqual(names, ["broken.png"])
|
||||
self.assertEqual(prompts["broken.png"], "")
|
||||
|
||||
def test_non_image_files_are_ignored(self):
|
||||
with open(os.path.join(self.temp_dir, "notes.txt"), "w", encoding="utf-8") as handle:
|
||||
handle.write("not an image")
|
||||
names, _, prompts = scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
self.assertEqual(names, [])
|
||||
self.assertEqual(prompts, {})
|
||||
|
||||
def test_corrupt_cache_file_is_ignored_and_replaced(self):
|
||||
self._make_png("image.png", 1_700_000_000)
|
||||
with open(self.cache_file, "w", encoding="utf-8") as handle:
|
||||
handle.write("{ broken json")
|
||||
|
||||
names, _, _ = scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
self.assertEqual(names, ["image.png"])
|
||||
with open(self.cache_file, "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
self.assertIsInstance(payload, dict)
|
||||
self.assertEqual(len(payload), 1)
|
||||
|
||||
def test_cache_is_pruned_to_256_entries(self):
|
||||
self._make_png("image.png", 1_700_000_000)
|
||||
stale = {
|
||||
os.path.join(self.temp_dir, f"stale-{index}.png"): {
|
||||
"mtime": 1,
|
||||
"prompt": f"old-{index}",
|
||||
}
|
||||
for index in range(300)
|
||||
}
|
||||
with open(self.cache_file, "w", encoding="utf-8") as handle:
|
||||
json.dump(stale, handle)
|
||||
|
||||
scan_image_folder(self.temp_dir, metadata_limit=10)
|
||||
with open(self.cache_file, "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
self.assertEqual(len(payload), 256)
|
||||
self.assertIn(os.path.join(self.temp_dir, "image.png"), payload)
|
||||
|
||||
def test_byte_parameters_and_positive_parser(self):
|
||||
self.assertEqual(
|
||||
parse_positive_from_parameters(b"A cat\r\nNegative prompt: dog\r\nSteps: 20"),
|
||||
"A cat",
|
||||
)
|
||||
self.assertEqual(parse_positive_from_parameters(None), "")
|
||||
self.assertEqual(parse_positive_from_parameters("Line one\nSteps: 2"), "Line one")
|
||||
self.assertEqual(decode_bytes_safely(b"plain"), "plain")
|
||||
self.assertEqual(decode_bytes_safely(None), "")
|
||||
|
||||
def test_graph_tracing_handles_cycles_and_ignores_negative_links(self):
|
||||
graph = {
|
||||
"1": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["2", 0], "negative": ["4", 0]},
|
||||
},
|
||||
"2": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "hello", "positive": ["3", 0]},
|
||||
},
|
||||
"3": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "world", "positive": ["2", 0]},
|
||||
},
|
||||
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": "negative text"}},
|
||||
}
|
||||
self.assertEqual(trace_positive_prompt_from_graph(graph), "hello\nworld")
|
||||
|
||||
fallback = {"7": {"class_type": "CLIPTextEncode", "inputs": {"text": "fallback text"}}}
|
||||
self.assertEqual(trace_positive_prompt_from_graph(fallback), "fallback text")
|
||||
self.assertEqual(trace_positive_prompt_from_graph(None), "")
|
||||
self.assertEqual(trace_positive_prompt_from_graph({}), "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,43 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import tempfile
|
||||
import shutil
|
||||
from PIL import Image
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.load_recent import LoadRecentOutputs
|
||||
|
||||
class TestLoadRecent(unittest.TestCase):
|
||||
def setUp(self):
|
||||
import folder_paths
|
||||
self.node = LoadRecentOutputs()
|
||||
out_dir = folder_paths.get_output_directory()
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
self.temp_dir = tempfile.mkdtemp(dir=out_dir)
|
||||
self.rel_dir = os.path.relpath(self.temp_dir, out_dir)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_load_single_image_finds_recent_image(self):
|
||||
img = Image.new("RGB", (64, 64), color="red")
|
||||
f1 = os.path.join(self.temp_dir, "recent_01.png")
|
||||
f2 = os.path.join(self.temp_dir, "recent_02.png")
|
||||
img.save(f1)
|
||||
img.save(f2)
|
||||
|
||||
out_tensor, = self.node.load_single_image(output_folder=self.rel_dir, amount=2, index=0)
|
||||
self.assertEqual(out_tensor.shape[0], 1)
|
||||
self.assertEqual(out_tensor.shape[1], 64)
|
||||
self.assertEqual(out_tensor.shape[2], 64)
|
||||
|
||||
def test_load_recent_empty_directory_returns_dummy(self):
|
||||
out_tensor, = self.node.load_single_image(output_folder=self.rel_dir, amount=5, index=0)
|
||||
self.assertEqual(out_tensor.shape, (1, 512, 512, 3))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,144 @@
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper # noqa: F401
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from nodes.load_recent import LoadRecentOutputs
|
||||
|
||||
|
||||
class TestLoadRecentExtended(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = LoadRecentOutputs()
|
||||
out_dir = folder_paths.get_output_directory()
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
self.temp_dir = tempfile.mkdtemp(dir=out_dir)
|
||||
self.rel_dir = os.path.relpath(self.temp_dir, out_dir)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def _make_image(self, name, color, mtime):
|
||||
path = os.path.join(self.temp_dir, name)
|
||||
Image.new("RGB", (4, 4), color=color).save(path)
|
||||
os.utime(path, (mtime, mtime))
|
||||
return path
|
||||
|
||||
def _load(self, amount=5, index=0, order="Newest First (0 = latest)", folder=None):
|
||||
tensor, = self.node.load_single_image(
|
||||
output_folder=self.rel_dir if folder is None else folder,
|
||||
amount=amount,
|
||||
index=index,
|
||||
order=order,
|
||||
)
|
||||
return tensor
|
||||
|
||||
def test_equal_mtimes_have_a_deterministic_tiebreak(self):
|
||||
base = 1_700_000_000
|
||||
self._make_image("z.png", "red", base)
|
||||
self._make_image("a.png", "green", base)
|
||||
self._make_image("m.png", "blue", base)
|
||||
|
||||
# Newest-first with equal mtimes sorts by full path descending.
|
||||
self.assertEqual(tuple(self._load(amount=3, index=0).shape), (1, 4, 4, 3))
|
||||
|
||||
colors = [self._load(amount=3, index=i)[0, 0, 0].tolist() for i in range(3)]
|
||||
self.assertEqual(colors[0][0], 1.0) # z.png red
|
||||
self.assertEqual(colors[1][2], 1.0) # m.png blue
|
||||
self.assertAlmostEqual(colors[2][1], 0.5019608, places=5) # a.png green
|
||||
|
||||
def test_chronological_order_reverses_the_same_set(self):
|
||||
base = 1_700_000_000
|
||||
self._make_image("oldest.png", "red", base)
|
||||
self._make_image("middle.png", "green", base + 10)
|
||||
self._make_image("newest.png", "blue", base + 20)
|
||||
|
||||
newest_first = [self._load(amount=3, index=i)[0, 0, 0].tolist() for i in range(3)]
|
||||
oldest_first = [
|
||||
self._load(amount=3, index=i, order="Chronological (Oldest First)")[0, 0, 0].tolist()
|
||||
for i in range(3)
|
||||
]
|
||||
self.assertEqual(oldest_first, list(reversed(newest_first)))
|
||||
|
||||
def test_index_wraps_and_negative_index_clamps_to_zero(self):
|
||||
base = 1_700_000_000
|
||||
self._make_image("older.png", "red", base)
|
||||
self._make_image("newer.png", "blue", base + 10)
|
||||
|
||||
wrapped = self._load(amount=2, index=3)[0, 0, 0].tolist()
|
||||
self.assertEqual(wrapped[0], 1.0, "index 3 must wrap to the second of two images")
|
||||
|
||||
clamped = self._load(amount=2, index=-5)[0, 0, 0].tolist()
|
||||
self.assertEqual(clamped[2], 1.0, "a negative index must clamp to the newest image")
|
||||
|
||||
def test_invalid_amount_and_index_fall_back_safely(self):
|
||||
base = 1_700_000_000
|
||||
self._make_image("older.png", "red", base)
|
||||
self._make_image("newer.png", "blue", base + 10)
|
||||
|
||||
tensor = self._load(amount="not-a-number", index="not-a-number")
|
||||
self.assertEqual(tuple(tensor.shape), (1, 4, 4, 3))
|
||||
self.assertEqual(tensor[0, 0, 0].tolist()[2], 1.0)
|
||||
|
||||
tensor = self._load(amount=0, index=0)
|
||||
self.assertEqual(tensor[0, 0, 0].tolist()[2], 1.0)
|
||||
|
||||
def test_amount_larger_than_available_still_wraps(self):
|
||||
base = 1_700_000_000
|
||||
self._make_image("only.png", "green", base)
|
||||
tensor = self._load(amount=100, index=7)
|
||||
self.assertEqual(tuple(tensor.shape), (1, 4, 4, 3))
|
||||
self.assertAlmostEqual(float(tensor[0, 0, 0, 1]), 0.5019608, places=5)
|
||||
|
||||
def test_missing_or_empty_folder_returns_dummy(self):
|
||||
missing = os.path.join(self.temp_dir, "does-not-exist")
|
||||
tensor = self._load(folder=os.path.relpath(missing, folder_paths.get_output_directory()))
|
||||
self.assertEqual(tuple(tensor.shape), (1, 512, 512, 3))
|
||||
self.assertEqual(tensor.dtype, torch.float32)
|
||||
self.assertEqual(float(tensor.sum()), 0.0)
|
||||
|
||||
empty = os.path.join(self.temp_dir, "empty")
|
||||
os.makedirs(empty)
|
||||
tensor = self._load(folder=os.path.relpath(empty, folder_paths.get_output_directory()))
|
||||
self.assertEqual(tuple(tensor.shape), (1, 512, 512, 3))
|
||||
self.assertEqual(float(tensor.sum()), 0.0)
|
||||
|
||||
def test_corrupt_image_returns_dummy_instead_of_raising(self):
|
||||
path = os.path.join(self.temp_dir, "broken.png")
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(b"this is not a PNG")
|
||||
os.utime(path, (1_700_000_000, 1_700_000_000))
|
||||
|
||||
tensor = self._load(amount=1, index=0)
|
||||
self.assertEqual(tuple(tensor.shape), (1, 512, 512, 3))
|
||||
self.assertEqual(float(tensor.sum()), 0.0)
|
||||
|
||||
def test_unsupported_extensions_are_ignored(self):
|
||||
with open(os.path.join(self.temp_dir, "notes.txt"), "w", encoding="utf-8") as handle:
|
||||
handle.write("not an image")
|
||||
tensor = self._load(amount=5, index=0)
|
||||
self.assertEqual(tuple(tensor.shape), (1, 512, 512, 3))
|
||||
self.assertEqual(float(tensor.sum()), 0.0)
|
||||
|
||||
def test_returned_tensor_is_float_rgb_in_zero_one_range(self):
|
||||
path = os.path.join(self.temp_dir, "solid.png")
|
||||
Image.new("RGB", (3, 2), color=(255, 0, 0)).save(path)
|
||||
os.utime(path, (1_700_000_000, 1_700_000_000))
|
||||
|
||||
tensor = self._load(amount=1, index=0)
|
||||
self.assertEqual(tuple(tensor.shape), (1, 2, 3, 3))
|
||||
self.assertEqual(tensor.dtype, torch.float32)
|
||||
self.assertEqual(tensor[0, :, :, 0].tolist(), [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
|
||||
self.assertEqual(tensor[0, :, :, 1].tolist(), [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,185 @@
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
from nodes.local_runner import (
|
||||
RunLocalFileNode,
|
||||
get_allowed_scripts_directory,
|
||||
is_script_path_permitted,
|
||||
resolve_target_path,
|
||||
authorize_node_run,
|
||||
consume_node_authorization,
|
||||
clear_all_authorizations
|
||||
)
|
||||
|
||||
class TestLocalRunner(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = RunLocalFileNode()
|
||||
clear_all_authorizations()
|
||||
|
||||
def tearDown(self):
|
||||
clear_all_authorizations()
|
||||
|
||||
def test_operator_authorization_required(self):
|
||||
with patch("nodes.local_runner.is_local_file_execution_enabled", return_value=True):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
script_path = os.path.join(temp_dir, "test.bat")
|
||||
with open(script_path, "w") as f:
|
||||
f.write("@echo off\n")
|
||||
|
||||
with patch("nodes.local_runner.get_allowed_scripts_directory", return_value=temp_dir):
|
||||
stdout, stderr, exit_code, success, passthrough = self.node.run_local_file(
|
||||
file_path="test.bat",
|
||||
unique_id="node_99",
|
||||
trigger="trigger_data"
|
||||
)
|
||||
self.assertEqual(exit_code, -1)
|
||||
self.assertFalse(success)
|
||||
self.assertIn("Authorize Run", stderr)
|
||||
self.assertEqual(passthrough, "trigger_data")
|
||||
|
||||
def test_global_security_setting_disabled_by_default(self):
|
||||
with patch("nodes.local_runner.is_local_file_execution_enabled", return_value=False):
|
||||
stdout, stderr, exit_code, success, _ = self.node.run_local_file(
|
||||
file_path="some_script.bat",
|
||||
unique_id="node_99"
|
||||
)
|
||||
self.assertEqual(exit_code, -1)
|
||||
self.assertFalse(success)
|
||||
self.assertIn("disabled by default for security", stderr)
|
||||
|
||||
def test_path_confinement_rejects_absolute_paths(self):
|
||||
resolved, msg = resolve_target_path("C:\\Windows\\System32\\cmd.exe")
|
||||
self.assertIsNone(resolved)
|
||||
self.assertIn("Absolute paths are forbidden", msg)
|
||||
|
||||
resolved_posix, msg_posix = resolve_target_path("/bin/sh")
|
||||
self.assertIsNone(resolved_posix)
|
||||
self.assertIn("Absolute paths are forbidden", msg_posix)
|
||||
|
||||
def test_path_confinement_rejects_traversal(self):
|
||||
resolved, msg = resolve_target_path("../../../etc/passwd")
|
||||
self.assertIsNone(resolved)
|
||||
self.assertIn("Directory traversal", msg)
|
||||
|
||||
resolved_win, msg_win = resolve_target_path("scripts/../../sensitive.txt")
|
||||
self.assertIsNone(resolved_win)
|
||||
self.assertIn("Directory traversal", msg_win)
|
||||
|
||||
def test_single_use_authorization_consumed_on_run(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
if sys.platform.startswith("win"):
|
||||
script_path = os.path.join(temp_dir, "single_use.bat")
|
||||
with open(script_path, "w") as f:
|
||||
f.write("@echo off\necho Consumed\n")
|
||||
else:
|
||||
script_path = os.path.join(temp_dir, "single_use.sh")
|
||||
with open(script_path, "w") as f:
|
||||
f.write('#!/bin/bash\necho "Consumed"\n')
|
||||
|
||||
with patch("nodes.local_runner.get_allowed_scripts_directory", return_value=temp_dir), \
|
||||
patch("nodes.local_runner.is_local_file_execution_enabled", return_value=True):
|
||||
|
||||
# 1. Authorize for node_42
|
||||
filename = os.path.basename(script_path)
|
||||
authorize_node_run("node_42", filename)
|
||||
|
||||
# First run must succeed
|
||||
stdout, stderr, exit_code, success, _ = self.node.run_local_file(
|
||||
file_path=filename,
|
||||
unique_id="node_42"
|
||||
)
|
||||
self.assertEqual(exit_code, 0)
|
||||
self.assertTrue(success)
|
||||
|
||||
# Second run with same node ID must be rejected (authorization consumed)
|
||||
stdout2, stderr2, exit_code2, success2, _ = self.node.run_local_file(
|
||||
file_path=filename,
|
||||
unique_id="node_42"
|
||||
)
|
||||
self.assertEqual(exit_code2, -1)
|
||||
self.assertFalse(success2)
|
||||
self.assertIn("Node has not been authorized", stderr2)
|
||||
|
||||
def test_script_mismatch_authorization_rejected(self):
|
||||
with patch("nodes.local_runner.is_local_file_execution_enabled", return_value=True):
|
||||
authorize_node_run("node_1", "approved_tool.bat")
|
||||
authorized, msg = consume_node_authorization("node_1", "C:\\scripts\\malicious_tool.bat")
|
||||
self.assertFalse(authorized)
|
||||
self.assertIn("does not match", msg)
|
||||
|
||||
def test_build_command_args(self):
|
||||
target = "test_script.bat" if sys.platform.startswith("win") else "test_script.sh"
|
||||
cmd = self.node.build_command_args(target, '--arg1 "value with spaces" --num 42')
|
||||
if sys.platform.startswith("win"):
|
||||
self.assertIn("/c", cmd)
|
||||
self.assertIn("--arg1", cmd)
|
||||
self.assertTrue(any("value with spaces" in x for x in cmd))
|
||||
self.assertIn("42", cmd)
|
||||
else:
|
||||
self.assertIn("/bin/bash", cmd)
|
||||
self.assertIn("--arg1", cmd)
|
||||
self.assertIn("value with spaces", cmd)
|
||||
|
||||
def test_execution_synchronous(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
if sys.platform.startswith("win"):
|
||||
script_path = os.path.join(temp_dir, "echo_test.bat")
|
||||
with open(script_path, "w") as f:
|
||||
f.write("@echo off\necho Output: %1 %2\n")
|
||||
else:
|
||||
script_path = os.path.join(temp_dir, "echo_test.sh")
|
||||
with open(script_path, "w") as f:
|
||||
f.write('#!/bin/bash\necho "Output: $1 $2"\n')
|
||||
|
||||
filename = os.path.basename(script_path)
|
||||
authorize_node_run("node_sync", filename)
|
||||
|
||||
with patch("nodes.local_runner.get_allowed_scripts_directory", return_value=temp_dir), \
|
||||
patch("nodes.local_runner.is_local_file_execution_enabled", return_value=True):
|
||||
stdout, stderr, exit_code, success, passthrough = self.node.run_local_file(
|
||||
file_path=filename,
|
||||
parameters="Hello World",
|
||||
run_mode="Synchronous (Wait for Output)",
|
||||
trigger="flow_pass",
|
||||
unique_id="node_sync"
|
||||
)
|
||||
self.assertEqual(exit_code, 0)
|
||||
self.assertTrue(success)
|
||||
self.assertIn("Hello World", stdout)
|
||||
self.assertEqual(passthrough, "flow_pass")
|
||||
|
||||
def test_timeout_cancellation(self):
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
if sys.platform.startswith("win"):
|
||||
script_path = os.path.join(temp_dir, "sleep_test.bat")
|
||||
with open(script_path, "w") as f:
|
||||
f.write(f"@echo off\n\"{sys.executable}\" -c \"import time; time.sleep(5)\"\n")
|
||||
else:
|
||||
script_path = os.path.join(temp_dir, "sleep_test.sh")
|
||||
with open(script_path, "w") as f:
|
||||
f.write("#!/bin/bash\nsleep 3\n")
|
||||
|
||||
filename = os.path.basename(script_path)
|
||||
authorize_node_run("node_timeout", filename)
|
||||
|
||||
with patch("nodes.local_runner.get_allowed_scripts_directory", return_value=temp_dir), \
|
||||
patch("nodes.local_runner.is_local_file_execution_enabled", return_value=True):
|
||||
stdout, stderr, exit_code, success, _ = self.node.run_local_file(
|
||||
file_path=filename,
|
||||
timeout=1,
|
||||
run_mode="Synchronous (Wait for Output)",
|
||||
unique_id="node_timeout"
|
||||
)
|
||||
self.assertEqual(exit_code, -1)
|
||||
self.assertFalse(success)
|
||||
self.assertIn("timed out", stderr)
|
||||
import time
|
||||
time.sleep(0.5)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,58 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.lora_finder import TextLoraFinder
|
||||
|
||||
class TestLoraFinder(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.finder = TextLoraFinder()
|
||||
|
||||
def test_empty_or_none_text(self):
|
||||
model, clip, loras = self.finder.find_and_load_loras(
|
||||
model="MODEL_OBJ",
|
||||
clip="CLIP_OBJ",
|
||||
text=""
|
||||
)
|
||||
self.assertEqual(model, "MODEL_OBJ")
|
||||
self.assertEqual(clip, "CLIP_OBJ")
|
||||
self.assertEqual(loras, "")
|
||||
|
||||
m2, c2, l2 = self.finder.find_and_load_loras(
|
||||
model="MODEL_OBJ",
|
||||
clip="CLIP_OBJ",
|
||||
text=None
|
||||
)
|
||||
self.assertEqual(l2, "")
|
||||
|
||||
@patch("nodes.lora_finder.get_filtered_loras_mapping")
|
||||
@patch("nodes.lora_finder.folder_paths.get_full_path")
|
||||
@patch("nodes.lora_finder.os.path.exists")
|
||||
@patch("nodes.lora_finder.load_torch_file")
|
||||
@patch("nodes.lora_finder.comfy.sd.load_lora_for_models")
|
||||
def test_find_and_load_lora_match(self, mock_load_lora, mock_load_torch, mock_exists, mock_get_path, mock_mapping):
|
||||
mock_mapping.return_value = {
|
||||
"Cyberpunk Character": "characters/char_cyberpunk-character_v1.safetensors"
|
||||
}
|
||||
mock_get_path.return_value = "/models/loras/characters/char_cyberpunk-character_v1.safetensors"
|
||||
mock_exists.return_value = True
|
||||
mock_load_torch.return_value = {}
|
||||
mock_load_lora.return_value = ("PATCHED_MODEL", "PATCHED_CLIP")
|
||||
|
||||
model, clip, loras = self.finder.find_and_load_loras(
|
||||
model="MODEL",
|
||||
clip="CLIP",
|
||||
text="A digital painting of a Cyberpunk Character at neon midnight",
|
||||
search_for="Parsed Name",
|
||||
output_format="Parsed Name"
|
||||
)
|
||||
self.assertEqual(model, "PATCHED_MODEL")
|
||||
self.assertEqual(clip, "PATCHED_CLIP")
|
||||
self.assertIn("Cyberpunk Character", loras)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,202 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
import shutil
|
||||
import inspect
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
import nodes.lora_loader as lora_loader
|
||||
from nodes.lora_loader import (
|
||||
select_a_lora,
|
||||
save_lora_cycle_state,
|
||||
load_lora_cycle_state,
|
||||
increment_lora_usage,
|
||||
flush_lora_usage,
|
||||
load_usage_data,
|
||||
reset_lora_usage_file,
|
||||
get_file_sha256_cached,
|
||||
VisualLoraLoader,
|
||||
)
|
||||
|
||||
|
||||
class TestLoraSelection(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.old_state_file = lora_loader.LORA_STATE_FILE
|
||||
self.old_user_dir = lora_loader.USER_DIR
|
||||
self.old_usage_cache = lora_loader._USAGE_CACHE
|
||||
self.old_usage_dirty = lora_loader._USAGE_DIRTY
|
||||
self.old_usage_timer = lora_loader._USAGE_TIMER
|
||||
self.old_flush_delay = lora_loader._USAGE_FLUSH_DELAY
|
||||
self.old_hash_cache = lora_loader._HASH_CACHE
|
||||
self.old_hash_cache_file = lora_loader.HASH_CACHE_FILE
|
||||
|
||||
lora_loader.LORA_STATE_FILE = os.path.join(self.temp_dir, "lora_loader_state.json")
|
||||
lora_loader.HASH_CACHE_FILE = os.path.join(self.temp_dir, "lora_hash_cache.json")
|
||||
lora_loader.USER_DIR = self.temp_dir
|
||||
lora_loader._USAGE_CACHE = None
|
||||
lora_loader._USAGE_DIRTY = False
|
||||
lora_loader._USAGE_TIMER = None
|
||||
lora_loader._USAGE_FLUSH_DELAY = 60.0
|
||||
lora_loader._HASH_CACHE = None
|
||||
|
||||
def tearDown(self):
|
||||
if lora_loader._USAGE_TIMER is not None:
|
||||
lora_loader._USAGE_TIMER.cancel()
|
||||
lora_loader.LORA_STATE_FILE = self.old_state_file
|
||||
lora_loader.USER_DIR = self.old_user_dir
|
||||
lora_loader._USAGE_CACHE = self.old_usage_cache
|
||||
lora_loader._USAGE_DIRTY = self.old_usage_dirty
|
||||
lora_loader._USAGE_TIMER = self.old_usage_timer
|
||||
lora_loader._USAGE_FLUSH_DELAY = self.old_flush_delay
|
||||
lora_loader._HASH_CACHE = self.old_hash_cache
|
||||
lora_loader.HASH_CACHE_FILE = self.old_hash_cache_file
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_sequential_advances_and_retry_does_not_advance(self):
|
||||
choices = ["A", "B", "C"]
|
||||
self.assertEqual(select_a_lora("Sequential", "node1", choices, "p1"), "A")
|
||||
self.assertEqual(select_a_lora("Sequential", "node1", choices, "p2"), "B")
|
||||
# Crash recovery / retry returns the assigned item without double-advancing.
|
||||
self.assertEqual(select_a_lora("Sequential", "node1", choices, "p1"), "A")
|
||||
self.assertEqual(select_a_lora("Sequential", "node1", choices, "p3"), "C")
|
||||
self.assertEqual(select_a_lora("Sequential", "node1", choices, "p4"), "A")
|
||||
|
||||
def test_sequential_reorder_keeps_cycle_position(self):
|
||||
# The state bucket is keyed by the selected set, so a different
|
||||
# click/toggle order must not shift the index onto another LoRA.
|
||||
self.assertEqual(select_a_lora("Sequential", "node_order", ["A", "B", "C"], "o1"), "A")
|
||||
self.assertEqual(select_a_lora("Sequential", "node_order", ["A", "B", "C"], "o2"), "B")
|
||||
|
||||
# Reordering after a reselect must continue the same cycle (C next),
|
||||
# never repeat B and never skip C.
|
||||
self.assertEqual(select_a_lora("Sequential", "node_order", ["C", "A", "B"], "o3"), "C")
|
||||
self.assertEqual(select_a_lora("Sequential", "node_order", ["B", "C", "A"], "o4"), "A")
|
||||
# Retrying a reordered prompt still replays its original assignment.
|
||||
self.assertEqual(select_a_lora("Sequential", "node_order", ["A", "C", "B"], "o3"), "C")
|
||||
|
||||
def test_sequential_changed_selection_starts_fresh_cycle(self):
|
||||
self.assertEqual(select_a_lora("Sequential", "node_set", ["A", "B"], "s1"), "A")
|
||||
self.assertEqual(select_a_lora("Sequential", "node_set", ["A", "B"], "s2"), "B")
|
||||
# A different selection set is a different state bucket.
|
||||
self.assertEqual(select_a_lora("Sequential", "node_set", ["A", "B", "C"], "s3"), "A")
|
||||
|
||||
def test_missing_selection_mode_defaults_to_sequential(self):
|
||||
params = inspect.signature(VisualLoraLoader.load_lora).parameters
|
||||
self.assertEqual(params["_selection_mode"].default, "Sequential")
|
||||
self.assertEqual(
|
||||
VisualLoraLoader.INPUT_TYPES()["hidden"]["_selection_mode"][1]["default"],
|
||||
"Sequential",
|
||||
)
|
||||
# IS_CHANGED must invalidate for the same fallback mode, otherwise a
|
||||
# prompt without _selection_mode would be treated as cacheable "All".
|
||||
self.assertIsInstance(VisualLoraLoader.IS_CHANGED(), float)
|
||||
|
||||
def test_no_replace_cycles_all_items_and_avoids_immediate_repeat(self):
|
||||
choices = ["A", "B", "C"]
|
||||
first_cycle = [
|
||||
select_a_lora("Random (No Replace)", "node2", choices, f"c1-{i}")
|
||||
for i in range(3)
|
||||
]
|
||||
self.assertEqual(set(first_cycle), {"A", "B", "C"})
|
||||
|
||||
second_cycle = [
|
||||
select_a_lora("Random (No Replace)", "node2", choices, f"c2-{i}")
|
||||
for i in range(3)
|
||||
]
|
||||
self.assertEqual(set(second_cycle), {"A", "B", "C"})
|
||||
# The first pick of a new cycle must not repeat the final pick of the
|
||||
# previous cycle back-to-back.
|
||||
self.assertNotEqual(second_cycle[0], first_cycle[-1])
|
||||
|
||||
def test_no_replace_retry_does_not_consume_pool(self):
|
||||
choices = ["A", "B", "C", "D"]
|
||||
first = select_a_lora("Random (No Replace)", "node3", choices, "retry-1")
|
||||
state_before = load_lora_cycle_state()
|
||||
key = list(state_before.keys())[0]
|
||||
pool_before = list(state_before[key]["pool"])
|
||||
|
||||
retried = select_a_lora("Random (No Replace)", "node3", choices, "retry-1")
|
||||
state_after = load_lora_cycle_state()
|
||||
self.assertEqual(first, retried)
|
||||
self.assertEqual(state_after[key]["pool"], pool_before)
|
||||
|
||||
def test_assigned_map_is_pruned_but_recent_retry_survives(self):
|
||||
choices = ["A", "B", "C"]
|
||||
for i in range(60):
|
||||
select_a_lora("Sequential", "node4", choices, f"prompt-{i}")
|
||||
|
||||
state = load_lora_cycle_state()
|
||||
key = list(state.keys())[0]
|
||||
assigned = state[key]["assigned"]
|
||||
self.assertLessEqual(len(assigned), 50)
|
||||
self.assertNotIn("prompt-0", assigned)
|
||||
self.assertIn("prompt-59", assigned)
|
||||
|
||||
# Retrying an old prompt that is no longer tracked simply picks a new
|
||||
# item instead of crashing, and the map stays bounded.
|
||||
select_a_lora("Sequential", "node4", choices, "prompt-0")
|
||||
state = load_lora_cycle_state()
|
||||
self.assertLessEqual(len(state[key]["assigned"]), 50)
|
||||
|
||||
def test_global_state_prune_keeps_recent_entries(self):
|
||||
state = {}
|
||||
for i in range(105):
|
||||
state[f"seq_node_{i}"] = {"index": 0, "assigned": {}, "last_updated": i}
|
||||
|
||||
save_lora_cycle_state(state)
|
||||
|
||||
self.assertEqual(len(state), 100)
|
||||
self.assertNotIn("seq_node_0", state)
|
||||
self.assertNotIn("seq_node_4", state)
|
||||
self.assertIn("seq_node_5", state)
|
||||
self.assertIn("seq_node_104", state)
|
||||
|
||||
def test_usage_increments_are_batched_until_flush(self):
|
||||
self.assertEqual(load_usage_data(), {})
|
||||
with patch("nodes.lora_loader.get_env_setting", return_value="true"):
|
||||
increment_lora_usage("folder/alpha_one.safetensors")
|
||||
increment_lora_usage("folder/alpha_one.safetensors")
|
||||
increment_lora_usage("folder/beta_two.safetensors")
|
||||
# Nothing written yet, but the in-memory view is already current.
|
||||
self.assertEqual(load_usage_data().get("One"), 2)
|
||||
flush_lora_usage()
|
||||
|
||||
usage_file = os.path.join(self.temp_dir, "lora_usage.json")
|
||||
self.assertTrue(os.path.exists(usage_file))
|
||||
with open(usage_file, "r", encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
self.assertEqual(raw.get("One"), 2)
|
||||
self.assertEqual(raw.get("Two"), 1)
|
||||
|
||||
reset_lora_usage_file()
|
||||
self.assertEqual(load_usage_data(), {})
|
||||
|
||||
def test_sha256_cache_reuses_unchanged_file(self):
|
||||
lora_file = os.path.join(self.temp_dir, "lora.safetensors")
|
||||
with open(lora_file, "wb") as f:
|
||||
f.write(b"weights-v1")
|
||||
|
||||
with patch("nodes.lora_loader.get_file_sha256", return_value="HASH-1") as mock_hash:
|
||||
first = get_file_sha256_cached(lora_file)
|
||||
second = get_file_sha256_cached(lora_file)
|
||||
self.assertEqual((first, second), ("HASH-1", "HASH-1"))
|
||||
self.assertEqual(mock_hash.call_count, 1)
|
||||
|
||||
# A changed mtime/size invalidates the cached hash.
|
||||
with open(lora_file, "ab") as f:
|
||||
f.write(b"-more")
|
||||
future = time.time() + 5
|
||||
os.utime(lora_file, (future, future))
|
||||
get_file_sha256_cached(lora_file)
|
||||
self.assertEqual(mock_hash.call_count, 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,147 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { stripImports } from "./fake_dom.mjs";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "persistent_queue.js");
|
||||
|
||||
let clock = 0;
|
||||
let timerId = 0;
|
||||
const timers = [];
|
||||
|
||||
function fakeSetTimeout(callback, delay = 0) {
|
||||
const timer = {
|
||||
id: ++timerId,
|
||||
at: clock + Math.max(0, Number(delay) || 0),
|
||||
callback,
|
||||
};
|
||||
timers.push(timer);
|
||||
return timer.id;
|
||||
}
|
||||
|
||||
function advance(ms) {
|
||||
const end = clock + Math.max(0, Number(ms) || 0);
|
||||
while (true) {
|
||||
let nextIndex = -1;
|
||||
for (let i = 0; i < timers.length; i++) {
|
||||
if (timers[i].at <= end && (nextIndex < 0 || timers[i].at < timers[nextIndex].at)) {
|
||||
nextIndex = i;
|
||||
}
|
||||
}
|
||||
if (nextIndex < 0) break;
|
||||
const [timer] = timers.splice(nextIndex, 1);
|
||||
clock = Math.max(clock, timer.at);
|
||||
timer.callback();
|
||||
}
|
||||
clock = end;
|
||||
}
|
||||
|
||||
async function flush() {
|
||||
for (let i = 0; i < 5; i++) await Promise.resolve();
|
||||
}
|
||||
|
||||
const apiEvents = new Map();
|
||||
const fetchCalls = [];
|
||||
let fetchImpl = async (url, options) => {
|
||||
fetchCalls.push({ url, options });
|
||||
return { ok: true };
|
||||
};
|
||||
|
||||
const api = {
|
||||
clientId: "client-abc",
|
||||
addEventListener(name, callback) {
|
||||
if (!apiEvents.has(name)) apiEvents.set(name, []);
|
||||
apiEvents.get(name).push(callback);
|
||||
},
|
||||
};
|
||||
|
||||
const extension = { name: "", setup: null };
|
||||
const app = {
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.setup = candidate.setup;
|
||||
},
|
||||
};
|
||||
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
app,
|
||||
api,
|
||||
authenticatedFetch: (url, options) => fetchImpl(url, options),
|
||||
setTimeout: fakeSetTimeout,
|
||||
clearTimeout: () => {},
|
||||
JSON,
|
||||
Promise,
|
||||
String,
|
||||
Object,
|
||||
Error,
|
||||
});
|
||||
|
||||
const source = stripImports(fs.readFileSync(sourcePath, "utf8"));
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
|
||||
assert.equal(typeof extension.setup, "function", "persistent queue extension must register");
|
||||
await extension.setup();
|
||||
|
||||
assert.deepEqual(
|
||||
[...apiEvents.keys()].sort(),
|
||||
["executing", "execution_start", "status"],
|
||||
"persistent queue must listen for queue status events"
|
||||
);
|
||||
|
||||
// No claim before the 500ms startup delay.
|
||||
advance(499);
|
||||
await flush();
|
||||
assert.equal(fetchCalls.length, 0, "queue ownership must not be claimed before startup settles");
|
||||
|
||||
advance(1);
|
||||
await flush();
|
||||
assert.equal(fetchCalls.length, 1, "queue ownership must be claimed after startup");
|
||||
assert.equal(fetchCalls[0].url, "/saturnnodes/persistent_queue/claim");
|
||||
assert.equal(fetchCalls[0].options.method, "POST");
|
||||
assert.equal(fetchCalls[0].options.headers["Content-Type"], "application/json");
|
||||
assert.deepEqual(
|
||||
JSON.parse(fetchCalls[0].options.body),
|
||||
{ client_id: "client-abc" },
|
||||
"claim body must carry the API client id"
|
||||
);
|
||||
|
||||
// Status events re-claim ownership (a reconnecting tab takes over recovery).
|
||||
apiEvents.get("status")[0]();
|
||||
await flush();
|
||||
assert.equal(fetchCalls.length, 2, "status events must re-claim queue ownership");
|
||||
assert.deepEqual(JSON.parse(fetchCalls[1].options.body), { client_id: "client-abc" });
|
||||
|
||||
// A client without an id must not send a claim.
|
||||
api.clientId = "";
|
||||
apiEvents.get("status")[0]();
|
||||
await flush();
|
||||
assert.equal(fetchCalls.length, 2, "clients without a client id must not claim ownership");
|
||||
|
||||
// A failed claim is contained and must not break later events.
|
||||
api.clientId = "client-xyz";
|
||||
fetchImpl = async () => {
|
||||
throw new Error("network unavailable");
|
||||
};
|
||||
apiEvents.get("status")[0]();
|
||||
await flush();
|
||||
assert.equal(fetchCalls.length, 2, "failed claim must not record a successful request");
|
||||
|
||||
// Restore a working transport and verify execution events do not throw when
|
||||
// rgthree-comfy is not installed (the dynamic import is caught internally).
|
||||
fetchImpl = async (url, options) => {
|
||||
fetchCalls.push({ url, options });
|
||||
return { ok: true };
|
||||
};
|
||||
apiEvents.get("execution_start")[0]({ detail: {} });
|
||||
apiEvents.get("executing")[0]({ detail: { prompt_id: "p1" } });
|
||||
await flush();
|
||||
advance(1200);
|
||||
await flush();
|
||||
|
||||
assert.equal(fetchCalls.length, 2, "execution events must not claim ownership");
|
||||
console.log("persistent queue claim DOM regression test passed");
|
||||
@@ -0,0 +1,381 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper # noqa: F401 (redirects SaturnNodes state to a temp dir)
|
||||
|
||||
import nodes.queue_control as queue_control
|
||||
from nodes.queue_control import PersistentQueueManager
|
||||
|
||||
|
||||
def make_entry(prompt_id, priority=0, item_pid=None, create_time=1, extra=None):
|
||||
item = [
|
||||
priority,
|
||||
prompt_id if item_pid is None else item_pid,
|
||||
{"prompt": {}},
|
||||
{"client_id": "old-client", "create_time": create_time},
|
||||
[],
|
||||
]
|
||||
entry = {"prompt_id": prompt_id, "item": item}
|
||||
if extra:
|
||||
entry.update(extra)
|
||||
return entry
|
||||
|
||||
|
||||
class FakePromptQueue:
|
||||
def __init__(self, history=None):
|
||||
self.history = {} if history is None else history
|
||||
self.items = []
|
||||
self.mutex = threading.RLock()
|
||||
self.currently_running = {}
|
||||
self.queue = []
|
||||
self.fail_delete = False
|
||||
self.fail_wipe = False
|
||||
self.fail_ids = set()
|
||||
|
||||
def get_history(self):
|
||||
return self.history
|
||||
|
||||
def put(self, item):
|
||||
if item[1] in self.fail_ids:
|
||||
raise RuntimeError(f"cannot queue {item[1]}")
|
||||
self.items.append(item)
|
||||
self.queue.append(item)
|
||||
|
||||
def task_done(self, item_id, history_result, status, process_item=None):
|
||||
self.currently_running.pop(item_id, None)
|
||||
|
||||
def delete_queue_item(self, predicate):
|
||||
if self.fail_delete:
|
||||
raise RuntimeError("delete failed")
|
||||
self.queue = [item for item in self.queue if not predicate(item)]
|
||||
|
||||
def wipe_queue(self):
|
||||
if self.fail_wipe:
|
||||
raise RuntimeError("wipe failed")
|
||||
self.queue = []
|
||||
|
||||
|
||||
class TestPersistentQueueRestore(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.TemporaryDirectory()
|
||||
self.temp_file = os.path.join(self.temp_dir.name, "persistent_queue.json")
|
||||
self.patches = [
|
||||
patch("nodes.queue_control.PERSISTENT_FILE", self.temp_file),
|
||||
patch("nodes.queue_control.USER_DIR", self.temp_dir.name),
|
||||
patch("nodes.queue_control.is_persistent_queue_enabled", return_value=True),
|
||||
patch.object(
|
||||
PersistentQueueManager,
|
||||
"_schedule_save_locked",
|
||||
new=lambda self: None,
|
||||
),
|
||||
]
|
||||
for patcher in self.patches:
|
||||
patcher.start()
|
||||
self.manager = PersistentQueueManager()
|
||||
|
||||
def tearDown(self):
|
||||
for patcher in reversed(self.patches):
|
||||
patcher.stop()
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
def _server(self, queue, number=0):
|
||||
return SimpleNamespace(
|
||||
prompt_queue=queue,
|
||||
number=number,
|
||||
queue_updated=MagicMock(),
|
||||
get_queue_status=MagicMock(return_value={}),
|
||||
send_sync=MagicMock(),
|
||||
)
|
||||
|
||||
def _write_payload(self, payload):
|
||||
with open(self.temp_file, "w", encoding="utf-8") as handle:
|
||||
json.dump(payload, handle)
|
||||
|
||||
def _read_payload(self):
|
||||
with open(self.temp_file, "r", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
|
||||
def test_normalize_entry_uses_top_level_prompt_id(self):
|
||||
raw = {"prompt_id": "top-level", "item": [0, "different", {"prompt": {}}]}
|
||||
entry = PersistentQueueManager._normalize_entry(raw)
|
||||
self.assertEqual(entry["prompt_id"], "top-level")
|
||||
self.assertEqual(entry["item"][1], "different")
|
||||
|
||||
self.assertIsNone(PersistentQueueManager._normalize_entry(None))
|
||||
self.assertIsNone(PersistentQueueManager._normalize_entry({"prompt_id": "x", "item": "bad"}))
|
||||
self.assertIsNone(PersistentQueueManager._normalize_entry({"prompt_id": "x", "item": [0]}))
|
||||
self.assertIsNone(PersistentQueueManager._normalize_entry({"item": [0, ""]}))
|
||||
self.assertEqual(
|
||||
PersistentQueueManager._normalize_entry({"item": [0, "from-item"]})["prompt_id"],
|
||||
"from-item",
|
||||
)
|
||||
|
||||
def test_normalize_items_repairs_and_dedupes_last_wins(self):
|
||||
first = make_entry("p1", priority=1)
|
||||
duplicate = make_entry("p1", priority=99)
|
||||
malformed = {"prompt_id": "bad", "item": "nope"}
|
||||
items, changed = PersistentQueueManager._normalize_items([first, malformed, duplicate])
|
||||
self.assertTrue(changed)
|
||||
self.assertEqual(len(items), 1)
|
||||
self.assertEqual(items[0]["prompt_id"], "p1")
|
||||
self.assertEqual(items[0]["item"][0], 99)
|
||||
|
||||
def test_restore_repairs_mismatched_item_prompt_id(self):
|
||||
entry = make_entry("top-level", item_pid="wrong-id", priority=4)
|
||||
self.manager.persistent_items = [entry]
|
||||
queue = FakePromptQueue()
|
||||
server = self._server(queue, number=2)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"):
|
||||
self.manager.restore_queue("new-client")
|
||||
|
||||
self.assertEqual(len(queue.items), 1)
|
||||
self.assertEqual(queue.items[0][1], "top-level")
|
||||
self.assertEqual(self.manager.persistent_items[0]["item"][1], "top-level")
|
||||
payload = self._read_payload()
|
||||
self.assertEqual(payload["items"][0]["item"][1], "top-level")
|
||||
|
||||
def test_restore_drops_prompts_already_in_history(self):
|
||||
self.manager.persistent_items = [make_entry("finished"), make_entry("pending", priority=2)]
|
||||
queue = FakePromptQueue(history={"finished": {"status": "success"}})
|
||||
server = self._server(queue)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"):
|
||||
self.manager.restore_queue()
|
||||
|
||||
self.assertEqual([item[1] for item in queue.items], ["pending"])
|
||||
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["pending"])
|
||||
|
||||
def test_restore_drops_malformed_entries(self):
|
||||
self.manager.persistent_items = [
|
||||
{"prompt_id": "short", "item": [0, "short"]},
|
||||
{"prompt_id": "bad-output", "item": [0, "bad-output", "not-a-dict"]},
|
||||
{"prompt_id": "valid", "item": [0, "valid", {"prompt": {}}, {}, []]},
|
||||
]
|
||||
queue = FakePromptQueue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"):
|
||||
self.manager.restore_queue()
|
||||
|
||||
self.assertEqual([item[1] for item in queue.items], ["valid"])
|
||||
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["valid"])
|
||||
|
||||
def test_restore_assigns_ordered_unique_create_times_and_client_id(self):
|
||||
self.manager.persistent_items = [
|
||||
make_entry("p1", priority=10),
|
||||
make_entry("p2", priority=5),
|
||||
make_entry("p3", priority=7),
|
||||
]
|
||||
queue = FakePromptQueue()
|
||||
server = self._server(queue, number=3)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"), \
|
||||
patch("nodes.queue_control.time.time", return_value=2_000_000.0):
|
||||
self.manager.restore_queue("client-42")
|
||||
|
||||
create_times = [item[3]["create_time"] for item in queue.items]
|
||||
self.assertEqual(create_times, [1_999_999_998, 1_999_999_999, 2_000_000_000])
|
||||
self.assertEqual(create_times, sorted(set(create_times)))
|
||||
for item in queue.items:
|
||||
self.assertEqual(item[3]["client_id"], "client-42")
|
||||
self.assertEqual(server.number, 11, "server priority counter must exceed restored items")
|
||||
|
||||
def test_restore_retries_with_malformed_retry_counter(self):
|
||||
entry = make_entry("flaky", extra={"retries": "not-a-number"})
|
||||
self.manager.persistent_items = [entry]
|
||||
queue = FakePromptQueue()
|
||||
queue.fail_ids.add("flaky")
|
||||
server = self._server(queue)
|
||||
|
||||
def env(key, default):
|
||||
if key == "PERSISTENT_QUEUE_MAX_RETRIES":
|
||||
return "2"
|
||||
return default
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"), \
|
||||
patch("nodes.queue_control.get_env_setting", side_effect=env):
|
||||
self.manager.restore_queue()
|
||||
self.assertEqual(self.manager.persistent_items[0]["retries"], 1)
|
||||
|
||||
self.manager.has_claimed_once = False
|
||||
self.manager.restore_queue()
|
||||
self.assertEqual(self.manager.persistent_items[0]["retries"], 2)
|
||||
|
||||
self.manager.has_claimed_once = False
|
||||
self.manager.restore_queue()
|
||||
|
||||
self.assertEqual(self.manager.persistent_items, [])
|
||||
self.assertEqual(queue.items, [])
|
||||
|
||||
def test_restore_drops_after_first_failure_when_retries_disabled(self):
|
||||
self.manager.persistent_items = [make_entry("flaky")]
|
||||
queue = FakePromptQueue()
|
||||
queue.fail_ids.add("flaky")
|
||||
server = self._server(queue)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"):
|
||||
self.manager.restore_queue()
|
||||
|
||||
self.assertEqual(self.manager.persistent_items, [])
|
||||
|
||||
def test_restore_tolerates_history_read_failure(self):
|
||||
self.manager.persistent_items = [make_entry("p1")]
|
||||
queue = FakePromptQueue()
|
||||
|
||||
def broken_history():
|
||||
raise RuntimeError("history unavailable")
|
||||
|
||||
queue.get_history = broken_history
|
||||
server = self._server(queue)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"):
|
||||
self.manager.restore_queue()
|
||||
|
||||
self.assertEqual([item[1] for item in queue.items], ["p1"])
|
||||
|
||||
def test_restore_force_state_setting_is_case_insensitive(self):
|
||||
for setting, expected in (("force paused", True), ("FORCE RUNNING", False)):
|
||||
self.manager.persistent_items = [make_entry("p1")]
|
||||
self.manager.has_claimed_once = False
|
||||
queue = FakePromptQueue()
|
||||
server = self._server(queue)
|
||||
|
||||
def env(key, default, value=setting):
|
||||
if key == "PERSISTENT_QUEUE_RESTORED_STATE":
|
||||
return value
|
||||
return default
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause") as set_pause, \
|
||||
patch("nodes.queue_control.get_env_setting", side_effect=env):
|
||||
self.manager.restore_queue()
|
||||
|
||||
set_pause.assert_called_once_with(expected)
|
||||
|
||||
def test_restore_is_skipped_when_disabled_or_empty(self):
|
||||
queue = FakePromptQueue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch("nodes.queue_control.is_persistent_queue_enabled", return_value=False):
|
||||
self.manager.persistent_items = [make_entry("p1")]
|
||||
self.manager.restore_queue()
|
||||
self.assertTrue(self.manager.has_claimed_once)
|
||||
self.assertEqual(queue.items, [])
|
||||
|
||||
self.manager.has_claimed_once = False
|
||||
self.manager.persistent_items = []
|
||||
with patch.object(queue_control.PromptServer, "instance", server):
|
||||
self.manager.restore_queue()
|
||||
self.assertEqual(queue.items, [])
|
||||
|
||||
def test_newer_schema_file_is_not_overwritten_by_restore_or_add(self):
|
||||
original = {
|
||||
"schema_version": 99,
|
||||
"items": [{"prompt_id": "future", "item": [0, "future", {}, {}, []], "future_field": True}],
|
||||
}
|
||||
self._write_payload(original)
|
||||
self.manager.load_from_file()
|
||||
self.assertTrue(self.manager._read_only_newer_schema)
|
||||
self.assertEqual(self.manager.persistent_items, [])
|
||||
|
||||
self.manager.add_item((0, "new", {"prompt": {}}, {}, []))
|
||||
self.manager.save_to_file(immediate=True)
|
||||
self.assertEqual(self._read_payload(), original)
|
||||
|
||||
def test_load_from_file_rewrites_duplicate_entries(self):
|
||||
duplicate = make_entry("p1", priority=99)
|
||||
self._write_payload({
|
||||
"schema_version": 1,
|
||||
"items": [make_entry("p1", priority=1), duplicate],
|
||||
})
|
||||
manager = PersistentQueueManager()
|
||||
self.assertEqual(len(manager.persistent_items), 1)
|
||||
self.assertEqual(manager.persistent_items[0]["item"][0], 99)
|
||||
self.assertEqual(len(self._read_payload()["items"]), 1)
|
||||
|
||||
def test_batch_info_requires_dict_and_ignores_finished_prompts(self):
|
||||
self.manager.set_batch_info("p1", "not-a-dict")
|
||||
self.manager.set_batch_info("", {"batchId": "empty"})
|
||||
self.assertEqual(self.manager.batch_meta, {})
|
||||
|
||||
info = {"batchId": "b1", "color": "#10b981"}
|
||||
self.manager.set_batch_info("p1", info)
|
||||
self.assertEqual(self.manager.batch_meta["p1"], info)
|
||||
|
||||
self.manager.add_item((0, "p1", {"prompt": {}}, {}, []))
|
||||
self.assertEqual(self.manager.persistent_items[0]["batch_info"], info)
|
||||
|
||||
self.manager.remove_item("p1")
|
||||
self.manager.set_batch_info("p1", {"batchId": "late"})
|
||||
self.assertNotIn("p1", self.manager.batch_meta)
|
||||
|
||||
def test_cancel_and_wipe_remove_persistence_only_after_success(self):
|
||||
queue = FakePromptQueue()
|
||||
server = self._server(queue)
|
||||
self.manager._patched = False
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server):
|
||||
self.manager.patch_server()
|
||||
|
||||
queue.put((0, "p1", {"prompt": {}}, {}, []))
|
||||
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p1"])
|
||||
|
||||
queue.fail_delete = True
|
||||
with self.assertRaises(RuntimeError):
|
||||
queue.delete_queue_item(lambda item: True)
|
||||
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p1"])
|
||||
|
||||
queue.fail_delete = False
|
||||
queue.delete_queue_item(lambda item: True)
|
||||
self.assertEqual(self.manager.persistent_items, [])
|
||||
|
||||
queue.put((0, "p2", {"prompt": {}}, {}, []))
|
||||
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p2"])
|
||||
queue.fail_wipe = True
|
||||
with self.assertRaises(RuntimeError):
|
||||
queue.wipe_queue()
|
||||
self.assertEqual([x["prompt_id"] for x in self.manager.persistent_items], ["p2"])
|
||||
|
||||
queue.fail_wipe = False
|
||||
queue.wipe_queue()
|
||||
self.assertEqual(self.manager.persistent_items, [])
|
||||
|
||||
queue.put((0, "p3", {"prompt": {}}, {}, []))
|
||||
queue.currently_running[7] = (0, "p3", {"prompt": {}}, {}, [])
|
||||
queue.task_done(7, {}, "success")
|
||||
self.assertEqual(self.manager.persistent_items, [])
|
||||
|
||||
def test_restore_syncs_client_id_to_all_entries(self):
|
||||
self.manager.persistent_items = [make_entry("p1"), make_entry("p2", priority=2)]
|
||||
queue = FakePromptQueue()
|
||||
server = self._server(queue)
|
||||
|
||||
with patch.object(queue_control.PromptServer, "instance", server), \
|
||||
patch.object(queue_control.pause_manager, "set_pause"):
|
||||
self.manager.restore_queue("client-x")
|
||||
|
||||
payload = self._read_payload()
|
||||
self.assertEqual(
|
||||
[item["item"][3]["client_id"] for item in payload["items"]],
|
||||
["client-x", "client-x"],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,117 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper # noqa: F401 (redirects SaturnNodes state to a temp dir)
|
||||
|
||||
import nodes.queue_control as queue_control
|
||||
from nodes.queue_control import PersistentQueueManager
|
||||
|
||||
ITEM = [0, "p1", {"prompt": {}}, {}, {}]
|
||||
|
||||
|
||||
class TestPersistentQueueSchema(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.TemporaryDirectory()
|
||||
self.temp_file = os.path.join(self.temp_dir.name, "persistent_queue.json")
|
||||
self.patches = [
|
||||
patch("nodes.queue_control.PERSISTENT_FILE", self.temp_file),
|
||||
patch("nodes.queue_control.USER_DIR", self.temp_dir.name),
|
||||
patch("nodes.queue_control.is_persistent_queue_enabled", return_value=True),
|
||||
]
|
||||
for patcher in self.patches:
|
||||
patcher.start()
|
||||
|
||||
def tearDown(self):
|
||||
for patcher in reversed(self.patches):
|
||||
patcher.stop()
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
def write_payload(self, payload):
|
||||
with open(self.temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(payload, f)
|
||||
|
||||
def read_payload(self):
|
||||
with open(self.temp_file, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
def test_legacy_list_is_migrated_to_versioned_envelope(self):
|
||||
self.write_payload([
|
||||
{
|
||||
"prompt_id": "p1",
|
||||
"item": ITEM,
|
||||
"batch_info": {"batchId": "b1", "color": "#4ade80"},
|
||||
}
|
||||
])
|
||||
|
||||
mgr = PersistentQueueManager()
|
||||
|
||||
self.assertEqual(len(mgr.persistent_items), 1)
|
||||
self.assertEqual(mgr.batch_meta["p1"]["batchId"], "b1")
|
||||
|
||||
payload = self.read_payload()
|
||||
self.assertIsInstance(payload, dict)
|
||||
self.assertEqual(
|
||||
payload["schema_version"],
|
||||
queue_control.PERSISTENT_QUEUE_SCHEMA_VERSION,
|
||||
)
|
||||
self.assertEqual(len(payload["items"]), 1)
|
||||
self.assertEqual(payload["items"][0]["prompt_id"], "p1")
|
||||
|
||||
def test_current_envelope_loads_without_rewriting(self):
|
||||
original = {
|
||||
"schema_version": queue_control.PERSISTENT_QUEUE_SCHEMA_VERSION,
|
||||
"items": [{"prompt_id": "p2", "item": ITEM}],
|
||||
}
|
||||
self.write_payload(original)
|
||||
|
||||
mgr = PersistentQueueManager()
|
||||
|
||||
self.assertEqual(len(mgr.persistent_items), 1)
|
||||
self.assertEqual(mgr.persistent_items[0]["prompt_id"], "p2")
|
||||
self.assertEqual(self.read_payload(), original)
|
||||
|
||||
def test_newer_schema_is_left_untouched(self):
|
||||
original = {
|
||||
"schema_version": queue_control.PERSISTENT_QUEUE_SCHEMA_VERSION + 999,
|
||||
"items": [{"prompt_id": "p3", "item": ITEM}],
|
||||
}
|
||||
self.write_payload(original)
|
||||
|
||||
mgr = PersistentQueueManager()
|
||||
|
||||
self.assertEqual(mgr.persistent_items, [])
|
||||
self.assertEqual(
|
||||
self.read_payload(),
|
||||
original,
|
||||
"a file from a newer version must not be overwritten on downgrade",
|
||||
)
|
||||
|
||||
# Even an explicit save must keep the foreign file read-only.
|
||||
mgr.add_item((0, "p5", {}, {}, {}))
|
||||
mgr.flush()
|
||||
self.assertEqual(
|
||||
self.read_payload(),
|
||||
original,
|
||||
"saving while downgraded must not overwrite the newer file",
|
||||
)
|
||||
|
||||
def test_save_writes_versioned_envelope(self):
|
||||
mgr = PersistentQueueManager()
|
||||
mgr.add_item((0, "p4", {}, {}, {}))
|
||||
mgr.flush()
|
||||
|
||||
payload = self.read_payload()
|
||||
self.assertEqual(
|
||||
payload["schema_version"],
|
||||
queue_control.PERSISTENT_QUEUE_SCHEMA_VERSION,
|
||||
)
|
||||
self.assertEqual([x["prompt_id"] for x in payload["items"]], ["p4"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,123 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.queue_control import (
|
||||
PowerControlManager,
|
||||
PauseQueueManager,
|
||||
_issue_power_ticket,
|
||||
_consume_power_ticket,
|
||||
_POWER_TICKETS,
|
||||
_POWER_TICKET_LOCK
|
||||
)
|
||||
|
||||
class TestPowerControlManager(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.pause_manager = PauseQueueManager()
|
||||
self.pause_manager.paused = False
|
||||
self.pause_manager.is_waiting = False
|
||||
self.power_manager = PowerControlManager(self.pause_manager)
|
||||
with _POWER_TICKET_LOCK:
|
||||
_POWER_TICKETS.clear()
|
||||
|
||||
def tearDown(self):
|
||||
with _POWER_TICKET_LOCK:
|
||||
_POWER_TICKETS.clear()
|
||||
|
||||
def test_arm_disabled_by_default(self):
|
||||
with patch("nodes.queue_control.is_process_management_enabled", return_value=False):
|
||||
result = self.power_manager.arm("restart")
|
||||
self.assertFalse(result)
|
||||
self.assertIsNone(self.power_manager.pending_action)
|
||||
self.assertIsNone(self.power_manager.armed_at)
|
||||
status = self.power_manager.get_status()
|
||||
self.assertFalse(status["enabled"])
|
||||
self.assertFalse(self.power_manager.execute_restart())
|
||||
self.assertFalse(self.power_manager.execute_shutdown())
|
||||
|
||||
def test_arm_and_disarm(self):
|
||||
with patch("nodes.queue_control.is_process_management_enabled", return_value=True):
|
||||
result = self.power_manager.arm("restart")
|
||||
self.assertTrue(result)
|
||||
self.assertEqual(self.power_manager.pending_action, "restart")
|
||||
self.assertIsNotNone(self.power_manager.armed_at)
|
||||
|
||||
status = self.power_manager.get_status()
|
||||
self.assertEqual(status["pending_action"], "restart")
|
||||
self.assertTrue(status["enabled"])
|
||||
|
||||
self.power_manager.arm("shutdown")
|
||||
self.assertEqual(self.power_manager.pending_action, "shutdown")
|
||||
|
||||
self.power_manager.arm(None)
|
||||
self.assertIsNone(self.power_manager.pending_action)
|
||||
self.assertIsNone(self.power_manager.armed_at)
|
||||
|
||||
def test_power_ticket_single_use_handshake(self):
|
||||
# 1. Issue ticket for restart
|
||||
ticket = _issue_power_ticket("restart")
|
||||
self.assertIsInstance(ticket, str)
|
||||
self.assertGreater(len(ticket), 16)
|
||||
|
||||
# 2. Consume ticket with wrong action -> Rejected
|
||||
self.assertFalse(_consume_power_ticket(ticket, "shutdown"))
|
||||
|
||||
# 3. Consume ticket with correct action -> Accepted
|
||||
self.assertTrue(_consume_power_ticket(ticket, "restart"))
|
||||
|
||||
# 4. Consume ticket again -> Rejected (single-use)
|
||||
self.assertFalse(_consume_power_ticket(ticket, "restart"))
|
||||
|
||||
def test_power_ticket_expiration(self):
|
||||
ticket = _issue_power_ticket("shutdown")
|
||||
# Artificially expire the ticket
|
||||
with _POWER_TICKET_LOCK:
|
||||
_POWER_TICKETS[ticket] = (time.time() - 1.0, "shutdown")
|
||||
|
||||
self.assertFalse(_consume_power_ticket(ticket, "shutdown"))
|
||||
|
||||
def test_does_not_execute_when_paused(self):
|
||||
with patch("nodes.queue_control.is_process_management_enabled", return_value=True):
|
||||
self.power_manager.arm("restart")
|
||||
self.pause_manager.paused = True
|
||||
|
||||
with patch.object(self.power_manager, "execute_restart") as mock_restart:
|
||||
self.power_manager.check_and_execute()
|
||||
mock_restart.assert_not_called()
|
||||
self.assertEqual(self.power_manager.pending_action, "restart")
|
||||
|
||||
def test_does_not_execute_when_tasks_remaining(self):
|
||||
with patch("nodes.queue_control.is_process_management_enabled", return_value=True):
|
||||
self.power_manager.arm("restart")
|
||||
mock_queue = MagicMock()
|
||||
mock_queue.get_tasks_remaining.return_value = 2
|
||||
mock_queue.currently_running = {"1": True}
|
||||
|
||||
with patch("server.PromptServer.instance") as mock_server:
|
||||
mock_server.prompt_queue = mock_queue
|
||||
with patch.object(self.power_manager, "execute_restart") as mock_restart:
|
||||
self.power_manager.check_and_execute()
|
||||
mock_restart.assert_not_called()
|
||||
self.assertEqual(self.power_manager.pending_action, "restart")
|
||||
|
||||
def test_executes_when_queue_empty_and_idle(self):
|
||||
with patch("nodes.queue_control.is_process_management_enabled", return_value=True):
|
||||
self.power_manager.arm("restart")
|
||||
mock_queue = MagicMock()
|
||||
mock_queue.get_tasks_remaining.return_value = 0
|
||||
mock_queue.currently_running = {}
|
||||
|
||||
with patch("server.PromptServer.instance") as mock_server:
|
||||
mock_server.prompt_queue = mock_queue
|
||||
with patch.object(self.power_manager, "execute_restart") as mock_restart:
|
||||
self.power_manager.check_and_execute()
|
||||
mock_restart.assert_called_once()
|
||||
self.assertIsNone(self.power_manager.pending_action)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,143 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import {
|
||||
FakeMutationObserver,
|
||||
createDocument,
|
||||
createWindow,
|
||||
sleep,
|
||||
stripImports,
|
||||
} from "./fake_dom.mjs";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "js", "power_sidebar.js");
|
||||
|
||||
const document = createDocument();
|
||||
const window = createWindow(document);
|
||||
|
||||
const extension = { name: "", setup: null };
|
||||
const app = {
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.setup = candidate.setup;
|
||||
},
|
||||
extensionManager: {
|
||||
setting: { get: () => true },
|
||||
toast: { add() {} },
|
||||
},
|
||||
ui: { settings: { get: () => undefined, getSettingValue: () => undefined } },
|
||||
};
|
||||
|
||||
let source = stripImports(fs.readFileSync(sourcePath, "utf8"));
|
||||
source += `
|
||||
globalThis.__powerTest = {
|
||||
showPowerSidePopup,
|
||||
closePowerPopup,
|
||||
getActivePopup: () => activePopup,
|
||||
hasCleanup: () => typeof activePopupCleanup === "function",
|
||||
};
|
||||
`;
|
||||
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
document,
|
||||
window,
|
||||
app,
|
||||
api: { addEventListener() {}, fetchApi: async () => null },
|
||||
authenticatedFetch: async () => null,
|
||||
localStorage: { getItem: () => null, setItem() {} },
|
||||
MutationObserver: FakeMutationObserver,
|
||||
setTimeout,
|
||||
clearTimeout,
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
Date,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
});
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
|
||||
assert.equal(typeof extension.setup, "function", "power sidebar extension must register");
|
||||
|
||||
const { showPowerSidePopup, closePowerPopup, getActivePopup, hasCleanup } = context.__powerTest;
|
||||
|
||||
const anchor = document.createElement("button");
|
||||
anchor.id = "power-anchor";
|
||||
anchor.setRect({ left: 20, top: 850, right: 60, bottom: 890, width: 40, height: 40 });
|
||||
document.body.appendChild(anchor);
|
||||
|
||||
function pointerDown(target) {
|
||||
document.dispatchEvent({
|
||||
type: "pointerdown",
|
||||
target,
|
||||
preventDefault() {},
|
||||
stopPropagation() {},
|
||||
});
|
||||
}
|
||||
|
||||
// 1. Opening arms the outside-click listener exactly once.
|
||||
showPowerSidePopup(anchor);
|
||||
await sleep(25);
|
||||
const popupA = getActivePopup();
|
||||
assert.ok(popupA && popupA.isConnected, "first popup must be mounted");
|
||||
assert.equal(document.listenerCount("pointerdown"), 1, "one outside listener after opening");
|
||||
|
||||
// 2. Closing via a menu item must also remove its document listeners.
|
||||
closePowerPopup();
|
||||
assert.equal(getActivePopup(), null, "popup must be removed");
|
||||
assert.equal(document.listenerCount("pointerdown"), 0, "outside listener must be removed on close");
|
||||
assert.equal(document.listenerCount("keydown"), 0, "keydown listener must be removed on close");
|
||||
assert.equal(hasCleanup(), false, "cleanup reference must be cleared");
|
||||
|
||||
// 3. Reopen and click an item: the stale listener from popup A must not
|
||||
// close popup B before its item click runs.
|
||||
showPowerSidePopup(anchor);
|
||||
await sleep(25);
|
||||
const popupB = getActivePopup();
|
||||
assert.ok(popupB && popupB.isConnected, "second popup must be mounted");
|
||||
|
||||
const item = document.createElement("button");
|
||||
popupB.appendChild(item);
|
||||
pointerDown(item);
|
||||
assert.equal(getActivePopup(), popupB, "first click inside a reopened popup must keep it open");
|
||||
|
||||
// 4. A genuine outside click still closes it and cleans up.
|
||||
const outside = document.createElement("div");
|
||||
document.body.appendChild(outside);
|
||||
pointerDown(outside);
|
||||
assert.equal(getActivePopup(), null, "outside click must close the popup");
|
||||
assert.equal(document.listenerCount("pointerdown"), 0, "outside listener must be removed after closing");
|
||||
assert.equal(document.listenerCount("keydown"), 0, "keydown listener must be removed after closing");
|
||||
|
||||
// 5. Escape closes and cleans up too.
|
||||
showPowerSidePopup(anchor);
|
||||
await sleep(25);
|
||||
assert.equal(document.listenerCount("keydown"), 1, "keydown listener armed");
|
||||
document.dispatchEvent({ type: "keydown", key: "Escape", target: document.body });
|
||||
assert.equal(getActivePopup(), null, "Escape must close the popup");
|
||||
assert.equal(document.listenerCount("keydown"), 0, "Escape must remove the listener");
|
||||
|
||||
// 6. A popup detached by a DOM swap must not block reopening (toggle path).
|
||||
showPowerSidePopup(anchor);
|
||||
await sleep(25);
|
||||
const detached = getActivePopup();
|
||||
detached.remove();
|
||||
assert.equal(detached.isConnected, false, "popup is detached");
|
||||
showPowerSidePopup(anchor);
|
||||
await sleep(25);
|
||||
const reopened = getActivePopup();
|
||||
assert.ok(reopened && reopened.isConnected, "a detached popup must not swallow the reopen");
|
||||
closePowerPopup();
|
||||
|
||||
console.log("power popup DOM regression test passed");
|
||||
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.prompt_counter import PromptCounter
|
||||
|
||||
class TestPromptCounter(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = PromptCounter()
|
||||
|
||||
def test_empty_and_whitespace(self):
|
||||
txt, count = self.node.count_prompts("")
|
||||
self.assertEqual(txt, "")
|
||||
self.assertEqual(count, 0)
|
||||
|
||||
txt, count = self.node.count_prompts(" \n\n \t ")
|
||||
self.assertEqual(txt, " \n\n \t ")
|
||||
self.assertEqual(count, 0)
|
||||
|
||||
def test_single_prompt(self):
|
||||
sample = "A magnificent tree standing on a hill, photorealistic 8k"
|
||||
txt, count = self.node.count_prompts(sample)
|
||||
self.assertEqual(txt, sample)
|
||||
self.assertEqual(count, 1)
|
||||
|
||||
def test_newline_mode(self):
|
||||
sample = "Prompt 1\nPrompt 2\nPrompt 3\n\nPrompt 4"
|
||||
txt, count = self.node.count_prompts(sample, separator="Newline")
|
||||
self.assertEqual(txt, sample)
|
||||
self.assertEqual(count, 4)
|
||||
|
||||
def test_empty_line_mode(self):
|
||||
sample = "Line 1 of prompt 1\nLine 2 of prompt 1\n\nPrompt 2\n\nPrompt 3"
|
||||
txt, count = self.node.count_prompts(sample, separator=">1 Empty Line")
|
||||
self.assertEqual(txt, sample)
|
||||
self.assertEqual(count, 3)
|
||||
|
||||
def test_gt_two_empty_lines_mode(self):
|
||||
sample = "Prompt 1\n\nStill prompt 1\n\n\nPrompt 2"
|
||||
txt, count = self.node.count_prompts(sample, separator=">2 Empty Lines")
|
||||
self.assertEqual(txt, sample)
|
||||
self.assertEqual(count, 2)
|
||||
|
||||
def test_custom_regex(self):
|
||||
sample = "Prompt A\n---\nPrompt B\n---\nPrompt C"
|
||||
txt, count = self.node.count_prompts(sample, separator="Custom Regex", custom_regex=r"\n---\n")
|
||||
self.assertEqual(txt, sample)
|
||||
self.assertEqual(count, 3)
|
||||
|
||||
def test_custom_regex_capturing_group(self):
|
||||
sample = "Prompt A\n===\nPrompt B\n===\nPrompt C"
|
||||
txt, count = self.node.count_prompts(sample, separator="Custom Regex", custom_regex=r"(\n===\n)")
|
||||
self.assertEqual(txt, sample)
|
||||
self.assertEqual(count, 3)
|
||||
|
||||
def test_custom_regex_invalid_fallback(self):
|
||||
sample = "Prompt 1\n\nPrompt 2"
|
||||
txt, count = self.node.count_prompts(sample, separator="Custom Regex", custom_regex=r"[invalid-regex")
|
||||
self.assertEqual(txt, sample)
|
||||
self.assertEqual(count, 2)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,487 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import tempfile
|
||||
import json
|
||||
import shutil
|
||||
import time
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.prompt_iterator import (
|
||||
PromptQueueIterator,
|
||||
parse_prompt_blocks,
|
||||
load_state,
|
||||
save_state,
|
||||
reset_node_state,
|
||||
invalidate_state_cache,
|
||||
)
|
||||
import nodes.prompt_iterator as pi
|
||||
|
||||
class TestPromptQueueIterator(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = PromptQueueIterator()
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.state_file = os.path.join(self.temp_dir, "test_state.json")
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_parse_prompt_blocks_newline(self):
|
||||
text = "Prompt 1\nPrompt 2\n\nPrompt 3"
|
||||
blocks = parse_prompt_blocks(text, separator="Newline")
|
||||
self.assertEqual(blocks, ["Prompt 1", "Prompt 2", "Prompt 3"])
|
||||
|
||||
def test_parse_prompt_blocks_empty_lines(self):
|
||||
text = "Block 1 line 1\nBlock 1 line 2\n\nBlock 2"
|
||||
blocks = parse_prompt_blocks(text, separator=">1 Empty Line")
|
||||
self.assertEqual(len(blocks), 2)
|
||||
self.assertEqual(blocks[0], "Block 1 line 1\nBlock 1 line 2")
|
||||
self.assertEqual(blocks[1], "Block 2")
|
||||
|
||||
def test_parse_prompt_blocks_custom_regex(self):
|
||||
text = "Prompt A\n---\nPrompt B\n---\nPrompt C"
|
||||
blocks = parse_prompt_blocks(text, separator="Custom Regex", custom_regex=r"\n---\n")
|
||||
self.assertEqual(blocks, ["Prompt A", "Prompt B", "Prompt C"])
|
||||
|
||||
def test_parse_prompt_blocks_custom_regex_capturing_group(self):
|
||||
text = "Prompt A\n---\nPrompt B\n---\nPrompt C"
|
||||
blocks = parse_prompt_blocks(text, separator="Custom Regex", custom_regex=r"(\n---\n)")
|
||||
self.assertEqual(blocks, ["Prompt A", "Prompt B", "Prompt C"])
|
||||
|
||||
def test_parse_prompt_blocks_custom_regex_invalid_fallback(self):
|
||||
text = "Prompt 1\n\nPrompt 2"
|
||||
# Invalid regex [unterminated class
|
||||
blocks = parse_prompt_blocks(text, separator="Custom Regex", custom_regex=r"[unterminated")
|
||||
self.assertEqual(blocks, ["Prompt 1", "Prompt 2"])
|
||||
|
||||
def test_sequential_10_batch_execution(self):
|
||||
"""
|
||||
Verify that 10 prompts are executed in exact sequential order (0 to 9).
|
||||
"""
|
||||
prompts_10 = "\n\n".join([f"Prompt A{i}" for i in range(10)])
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
for expected_idx in range(10):
|
||||
p, rem_text, rem_count, exhausted = self.node.process_queue(
|
||||
pop_mode="Sequential (Loop on End)",
|
||||
separator=">1 Empty Line",
|
||||
text=prompts_10,
|
||||
unique_id="node_10_test"
|
||||
)
|
||||
self.assertEqual(p, f"Prompt A{expected_idx}")
|
||||
self.assertEqual(rem_count, 10 - (expected_idx + 1))
|
||||
|
||||
def test_linked_prompt_input_wins_over_stale_text(self):
|
||||
"""
|
||||
A connected `prompt` input is authoritative; the local multiline `text`
|
||||
widget must not replay stale characters from an earlier session.
|
||||
"""
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
stale_text = "stale one\n\nstale two"
|
||||
linked_text = "linked one\n\nlinked two"
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
p0, _, _, _ = self.node.process_queue(
|
||||
text=stale_text, prompt=linked_text, unique_id="linked_node"
|
||||
)
|
||||
p1, _, _, _ = self.node.process_queue(
|
||||
text=stale_text, prompt=linked_text, unique_id="linked_node"
|
||||
)
|
||||
|
||||
self.assertEqual(p0, "linked one")
|
||||
self.assertEqual(p1, "linked two")
|
||||
|
||||
def test_linked_empty_prompt_does_not_replay_local_text(self):
|
||||
"""An empty linked prompt produces no run, even if `text` is stale."""
|
||||
p, remaining, count, exhausted = self.node.process_queue(
|
||||
text="stale prompt text", prompt="", unique_id="linked_empty"
|
||||
)
|
||||
self.assertEqual(p, "")
|
||||
self.assertEqual(remaining, "")
|
||||
self.assertEqual(count, 0)
|
||||
self.assertFalse(exhausted)
|
||||
|
||||
def test_multi_pack_isolation_pack_a_and_pack_b(self):
|
||||
"""
|
||||
Verify Pack A (10 prompts) and Pack B (5 prompts) run with complete isolation.
|
||||
"""
|
||||
pack_a = "\n\n".join([f"PackA_{i}" for i in range(10)])
|
||||
pack_b = "\n\n".join([f"PackB_{i}" for i in range(5)])
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
# Run 3 items of Pack A
|
||||
p0, _, _, _ = self.node.process_queue(text=pack_a, unique_id="shared_node")
|
||||
p1, _, _, _ = self.node.process_queue(text=pack_a, unique_id="shared_node")
|
||||
p2, _, _, _ = self.node.process_queue(text=pack_a, unique_id="shared_node")
|
||||
self.assertEqual(p0, "PackA_0")
|
||||
self.assertEqual(p1, "PackA_1")
|
||||
self.assertEqual(p2, "PackA_2")
|
||||
|
||||
# Run Pack B -> Must start cleanly at PackB_0!
|
||||
pb0, _, _, _ = self.node.process_queue(text=pack_b, unique_id="shared_node")
|
||||
pb1, _, _, _ = self.node.process_queue(text=pack_b, unique_id="shared_node")
|
||||
self.assertEqual(pb0, "PackB_0")
|
||||
self.assertEqual(pb1, "PackB_1")
|
||||
|
||||
# Switch back to Pack A -> Must continue at PackA_3!
|
||||
p3, _, _, _ = self.node.process_queue(text=pack_a, unique_id="shared_node")
|
||||
self.assertEqual(p3, "PackA_3")
|
||||
|
||||
def test_text_modification_resets_counter_to_zero(self):
|
||||
"""
|
||||
Changing the prompt text (deleting 3 items) creates a new hash and starts at 0.
|
||||
"""
|
||||
pack_10 = "\n\n".join([f"Item_{i}" for i in range(10)])
|
||||
pack_7 = "\n\n".join([f"Item_{i}" for i in range(7)]) # 3 deleted
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
# Run 2 items of 10-pack
|
||||
self.node.process_queue(text=pack_10, unique_id="edit_node")
|
||||
self.node.process_queue(text=pack_10, unique_id="edit_node")
|
||||
|
||||
# Switch to edited 7-pack -> Must start at Item_0!
|
||||
first_of_7, _, rem_count, _ = self.node.process_queue(text=pack_7, unique_id="edit_node")
|
||||
self.assertEqual(first_of_7, "Item_0")
|
||||
self.assertEqual(rem_count, 6)
|
||||
|
||||
def test_forced_counter_reset(self):
|
||||
"""
|
||||
Manual counter reset forces next execution back to 0.
|
||||
"""
|
||||
pack = "P0\n\nP1\n\nP2\n\nP3"
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
# Run P0, P1
|
||||
p0, _, _, _ = self.node.process_queue(text=pack, unique_id="reset_test")
|
||||
p1, _, _, _ = self.node.process_queue(text=pack, unique_id="reset_test")
|
||||
self.assertEqual(p0, "P0")
|
||||
self.assertEqual(p1, "P1")
|
||||
|
||||
# Simulate Reset Counter (set index to 0)
|
||||
for k in list(state_db.keys()):
|
||||
if "reset_test" in k:
|
||||
state_db[k]["index"] = 0
|
||||
|
||||
# Next run must be P0 again!
|
||||
p0_again, _, _, _ = self.node.process_queue(text=pack, unique_id="reset_test")
|
||||
self.assertEqual(p0_again, "P0")
|
||||
|
||||
def test_crash_recovery_assigned_prompt_id_retains_index(self):
|
||||
"""
|
||||
Verify that if a prompt is interrupted/restored with the same prompt_id,
|
||||
it re-executes the exact same prompt index instead of skipping ahead.
|
||||
"""
|
||||
pack = "Prompt 0\n\nPrompt 1\n\nPrompt 2"
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
# 1. First execution with prompt_id = "job-uuid-1" -> gets Prompt 0
|
||||
with patch("nodes.prompt_iterator.get_current_prompt_id", return_value="job-uuid-1"):
|
||||
p0, _, _, _ = self.node.process_queue(text=pack, unique_id="recover_node")
|
||||
self.assertEqual(p0, "Prompt 0")
|
||||
|
||||
# 2. Workflow is interrupted mid-way and restored by persistent queue (same prompt_id "job-uuid-1")
|
||||
with patch("nodes.prompt_iterator.get_current_prompt_id", return_value="job-uuid-1"):
|
||||
p0_retry, _, _, _ = self.node.process_queue(text=pack, unique_id="recover_node")
|
||||
# Must STILL be Prompt 0, NOT Prompt 1!
|
||||
self.assertEqual(p0_retry, "Prompt 0")
|
||||
|
||||
# 3. Next workflow run with new prompt_id = "job-uuid-2" -> gets Prompt 1!
|
||||
with patch("nodes.prompt_iterator.get_current_prompt_id", return_value="job-uuid-2"):
|
||||
p1, _, _, _ = self.node.process_queue(text=pack, unique_id="recover_node")
|
||||
self.assertEqual(p1, "Prompt 1")
|
||||
|
||||
def test_per_node_hash_pruning(self):
|
||||
"""
|
||||
Verify that editing text in a node cleans up old abandoned text hashes beyond 3 entries.
|
||||
"""
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
# Create 6 text variants on the same node
|
||||
for i in range(6):
|
||||
text_variant = f"Variant {i} text"
|
||||
self.node.process_queue(text=text_variant, unique_id="prune_node")
|
||||
|
||||
# Only at most 3 entries for 'prune_node' should remain in state
|
||||
matching_keys = [k for k in state_db.keys() if k.startswith("node_prune_node_")]
|
||||
self.assertLessEqual(len(matching_keys), 3)
|
||||
|
||||
def test_random_cycle_permutation(self):
|
||||
"""
|
||||
Verify that Random (Cycle) mode visits all items in a permutation before repeating.
|
||||
"""
|
||||
pack = "A\n\nB\n\nC\n\nD\n\nE"
|
||||
state_db = {}
|
||||
|
||||
def mock_load():
|
||||
return dict(state_db)
|
||||
|
||||
def mock_save(new_state):
|
||||
nonlocal state_db
|
||||
state_db = dict(new_state)
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=mock_load), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=mock_save), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
picked = []
|
||||
for _ in range(5):
|
||||
p, _, _, _ = self.node.process_queue(
|
||||
pop_mode="Random (Cycle)",
|
||||
text=pack,
|
||||
unique_id="cycle_test_node"
|
||||
)
|
||||
picked.append(p)
|
||||
|
||||
# All 5 unique items must have been visited in 1 cycle
|
||||
self.assertEqual(len(picked), 5)
|
||||
self.assertEqual(set(picked), {"A", "B", "C", "D", "E"})
|
||||
|
||||
def test_separator_and_mode_isolate_state_keys(self):
|
||||
"""Changing the splitter/mode must not continue another config's pointer."""
|
||||
pack = "A\n\nB\n\nC"
|
||||
state_db = {}
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=lambda: dict(state_db)), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=lambda s: state_db.update(s)), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
a0, _, _, _ = self.node.process_queue(
|
||||
separator=">1 Empty Line", text=pack, unique_id="cfg_node")
|
||||
a1, _, _, _ = self.node.process_queue(
|
||||
separator=">1 Empty Line", text=pack, unique_id="cfg_node")
|
||||
self.assertEqual((a0, a1), ("A", "B"))
|
||||
|
||||
# Newline splitter is a different config and must start at index 0.
|
||||
n0, _, _, _ = self.node.process_queue(
|
||||
separator="Newline", text=pack, unique_id="cfg_node")
|
||||
self.assertEqual(n0, "A")
|
||||
|
||||
# Random mode with the same splitter is also a different config.
|
||||
r0, _, _, _ = self.node.process_queue(
|
||||
pop_mode="Random (Keep)", separator=">1 Empty Line", text=pack, unique_id="cfg_node")
|
||||
self.assertIn(r0, {"A", "B", "C"})
|
||||
|
||||
# Returning to the original config resumes at C.
|
||||
a2, _, _, _ = self.node.process_queue(
|
||||
separator=">1 Empty Line", text=pack, unique_id="cfg_node")
|
||||
self.assertEqual(a2, "C")
|
||||
|
||||
def test_input_types_preserve_legacy_widget_positions(self):
|
||||
# New widgets must be appended, never inserted in the middle: ComfyUI
|
||||
# restores workflow widget values positionally.
|
||||
spec = PromptQueueIterator.INPUT_TYPES()
|
||||
self.assertEqual(list(spec["required"].keys()), ["pop_mode", "separator", "text"])
|
||||
self.assertEqual(list(spec["optional"].keys())[:2], ["custom_regex", "track_by"])
|
||||
|
||||
def test_track_by_node_only_shares_pointer_across_text_edits(self):
|
||||
state_db = {}
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=lambda: dict(state_db)), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=lambda s: state_db.update(s)), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
p0, _, _, _ = self.node.process_queue(
|
||||
text="A\n\nB\n\nC", unique_id="node_track", track_by="Node Only")
|
||||
self.assertEqual(p0, "A")
|
||||
|
||||
# The edited prompt pack keeps the same pointer (index 1 -> Y).
|
||||
p1, _, _, _ = self.node.process_queue(
|
||||
text="X\n\nY\n\nZ", unique_id="node_track", track_by="Node Only")
|
||||
self.assertEqual(p1, "Y")
|
||||
|
||||
# Default tracking starts the edited pack from the beginning.
|
||||
d0, _, _, _ = self.node.process_queue(
|
||||
text="X\n\nY\n\nZ", unique_id="text_track")
|
||||
self.assertEqual(d0, "X")
|
||||
|
||||
def test_stop_on_end_exhausted_flag_and_remaining(self):
|
||||
state_db = {}
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=lambda: dict(state_db)), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=lambda s: state_db.update(s)), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
p0, rem0, count0, exhausted0 = self.node.process_queue(
|
||||
pop_mode="Sequential (Stop on End)", text="A\n\nB", unique_id="stop_node")
|
||||
self.assertEqual((p0, rem0, count0, exhausted0), ("A", "B", 1, False))
|
||||
|
||||
p1, rem1, count1, exhausted1 = self.node.process_queue(
|
||||
pop_mode="Sequential (Stop on End)", text="A\n\nB", unique_id="stop_node")
|
||||
self.assertEqual((p1, rem1, count1, exhausted1), ("B", "", 0, True))
|
||||
|
||||
# Further runs keep returning the final prompt and report exhaustion.
|
||||
p2, _, count2, exhausted2 = self.node.process_queue(
|
||||
pop_mode="Sequential (Stop on End)", text="A\n\nB", unique_id="stop_node")
|
||||
self.assertEqual((p2, count2, exhausted2), ("B", 0, True))
|
||||
|
||||
def test_random_cycle_remaining_follows_shuffle_order(self):
|
||||
state_db = {}
|
||||
|
||||
with patch("nodes.prompt_iterator.load_state", side_effect=lambda: dict(state_db)), \
|
||||
patch("nodes.prompt_iterator.save_state", side_effect=lambda s: state_db.update(s)), \
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"), \
|
||||
patch("nodes.prompt_iterator.random.shuffle", side_effect=lambda seq: seq.reverse()):
|
||||
results = []
|
||||
for _ in range(4):
|
||||
results.append(self.node.process_queue(
|
||||
pop_mode="Random (Cycle)", text="A\n\nB\n\nC\n\nD", unique_id="cycle_rem"))
|
||||
|
||||
# Reversed order is D, C, B, A.
|
||||
self.assertEqual([r[0] for r in results], ["D", "C", "B", "A"])
|
||||
self.assertEqual([r[2] for r in results], [3, 2, 1, 0])
|
||||
self.assertEqual(results[0][1], "C\n\nB\n\nA")
|
||||
self.assertEqual(results[1][1], "B\n\nA")
|
||||
self.assertEqual(results[2][1], "A")
|
||||
self.assertEqual(results[3][1], "")
|
||||
|
||||
def test_hard_reset_clears_assignments_and_shuffle(self):
|
||||
state = {
|
||||
"node_42_aaaa_1111": {
|
||||
"index": 3,
|
||||
"last_run": 3,
|
||||
"assigned_prompts": {"prompt-a": 1, "prompt-b": 2},
|
||||
"shuffled_order": [1, 0, 2],
|
||||
},
|
||||
"node_42": {"index": 2, "assigned_prompts": {"prompt-c": 0}},
|
||||
"node_99_bbbb_2222": {"index": 5, "assigned_prompts": {"keep": 1}},
|
||||
}
|
||||
|
||||
matched = reset_node_state(state, "42")
|
||||
|
||||
self.assertTrue(matched)
|
||||
self.assertEqual(state["node_42_aaaa_1111"]["index"], 0)
|
||||
self.assertEqual(state["node_42_aaaa_1111"]["last_run"], 0)
|
||||
self.assertEqual(state["node_42_aaaa_1111"]["assigned_prompts"], {})
|
||||
self.assertEqual(state["node_42_aaaa_1111"]["shuffled_order"], [])
|
||||
self.assertEqual(state["node_42"]["assigned_prompts"], {})
|
||||
# Other nodes are untouched.
|
||||
self.assertEqual(state["node_99_bbbb_2222"]["index"], 5)
|
||||
self.assertEqual(state["node_99_bbbb_2222"]["assigned_prompts"], {"keep": 1})
|
||||
|
||||
def test_global_prune_keeps_most_recent_entries(self):
|
||||
state = {}
|
||||
for i in range(105):
|
||||
state[f"node_prune_{i}_{i}"] = {"index": 0, "last_updated": i}
|
||||
|
||||
old_file = pi.STATE_FILE
|
||||
old_user = pi.USER_DIR
|
||||
pi.STATE_FILE = os.path.join(self.temp_dir, "prune_state.json")
|
||||
pi.USER_DIR = self.temp_dir
|
||||
try:
|
||||
save_state(state)
|
||||
finally:
|
||||
pi.STATE_FILE = old_file
|
||||
pi.USER_DIR = old_user
|
||||
invalidate_state_cache()
|
||||
|
||||
self.assertEqual(len(state), 100)
|
||||
self.assertNotIn("node_prune_0_0", state)
|
||||
self.assertNotIn("node_prune_4_4", state)
|
||||
self.assertIn("node_prune_5_5", state)
|
||||
self.assertIn("node_prune_104_104", state)
|
||||
|
||||
def test_load_state_uses_mtime_cache_and_returns_copies(self):
|
||||
state_file = os.path.join(self.temp_dir, "cache_state.json")
|
||||
old_file = pi.STATE_FILE
|
||||
old_user = pi.USER_DIR
|
||||
pi.STATE_FILE = state_file
|
||||
pi.USER_DIR = self.temp_dir
|
||||
try:
|
||||
invalidate_state_cache()
|
||||
save_state({"node_a": {"index": 1, "last_updated": 1}})
|
||||
first = load_state()
|
||||
first["node_a"]["index"] = 999 # Must not mutate the cache.
|
||||
second = load_state()
|
||||
self.assertEqual(second["node_a"]["index"], 1)
|
||||
|
||||
# External edit (new mtime) must invalidate the in-memory cache.
|
||||
with open(state_file, "w", encoding="utf-8") as f:
|
||||
json.dump({"node_a": {"index": 7, "last_updated": 2}}, f)
|
||||
future = time.time() + 5
|
||||
os.utime(state_file, (future, future))
|
||||
third = load_state()
|
||||
self.assertEqual(third["node_a"]["index"], 7)
|
||||
finally:
|
||||
pi.STATE_FILE = old_file
|
||||
pi.USER_DIR = old_user
|
||||
invalidate_state_cache()
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,123 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import {
|
||||
createDocument,
|
||||
createWindow,
|
||||
stripImports,
|
||||
} from "./fake_dom.mjs";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
const sourcePath = path.join(repoRoot, "web", "saturnnodes_settings.js");
|
||||
|
||||
const document = createDocument();
|
||||
const window = createWindow(document);
|
||||
|
||||
const extension = { name: "", setup: null };
|
||||
const app = {
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.setup = candidate.setup;
|
||||
},
|
||||
ui: { settings: { addSetting() {}, get: () => undefined, getSettingValue: () => undefined } },
|
||||
extensionManager: { toast: { add() {} } },
|
||||
graph: { setDirtyCanvas() {} },
|
||||
};
|
||||
|
||||
let source = stripImports(fs.readFileSync(sourcePath, "utf8"));
|
||||
source += "\nglobalThis.__settingsTest = { reapplyQueueHeightToMountedLists };\n";
|
||||
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
document,
|
||||
window,
|
||||
app,
|
||||
api: { addEventListener() {} },
|
||||
authenticatedFetch: async () => null,
|
||||
updateAllSaturnNodeColors() {},
|
||||
localStorage: { getItem: () => null, setItem() {} },
|
||||
fetch: async () => ({ ok: false, json: async () => ({}) }),
|
||||
MutationObserver: class { observe() {} disconnect() {} },
|
||||
setTimeout,
|
||||
clearTimeout,
|
||||
setInterval: () => 0,
|
||||
clearInterval: () => {},
|
||||
requestAnimationFrame: window.requestAnimationFrame,
|
||||
cancelAnimationFrame: window.cancelAnimationFrame,
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
Date,
|
||||
parseInt,
|
||||
parseFloat,
|
||||
isNaN,
|
||||
Infinity,
|
||||
});
|
||||
vm.runInContext(source, context, { filename: sourcePath });
|
||||
|
||||
assert.equal(typeof extension.setup, "function", "settings extension must register");
|
||||
const { reapplyQueueHeightToMountedLists } = context.__settingsTest;
|
||||
|
||||
const overlay = document.createElement("div");
|
||||
overlay.setAttribute("data-testid", "queue-progress-overlay");
|
||||
document.body.appendChild(overlay);
|
||||
|
||||
const firstList = document.createElement("div");
|
||||
firstList.className = "comfy-list";
|
||||
overlay.appendChild(firstList);
|
||||
|
||||
const applied = [];
|
||||
const applyHeight = height => {
|
||||
applied.push(height);
|
||||
// Mirrors applyQueueOverlayHeight(): every processed list gets the marker.
|
||||
for (const list of overlay.querySelectorAll(".comfy-list")) {
|
||||
list.dataset.saturnQueueHeight = String(height);
|
||||
}
|
||||
};
|
||||
|
||||
assert.equal(
|
||||
reapplyQueueHeightToMountedLists(overlay, 360, applyHeight),
|
||||
true,
|
||||
"a freshly mounted list must trigger a re-apply"
|
||||
);
|
||||
assert.deepEqual(applied, [360], "the saved height must be applied");
|
||||
|
||||
assert.equal(
|
||||
reapplyQueueHeightToMountedLists(overlay, 360, applyHeight),
|
||||
false,
|
||||
"an already-styled list must not be rewritten on every mutation"
|
||||
);
|
||||
assert.deepEqual(applied, [360], "no duplicate apply");
|
||||
|
||||
const secondList = document.createElement("div");
|
||||
secondList.className = "comfy-list";
|
||||
overlay.appendChild(secondList);
|
||||
|
||||
assert.equal(
|
||||
reapplyQueueHeightToMountedLists(overlay, 360, applyHeight),
|
||||
true,
|
||||
"a second list mounting later must also be styled"
|
||||
);
|
||||
assert.deepEqual(applied, [360, 360], "re-apply covers the new list");
|
||||
|
||||
assert.equal(
|
||||
reapplyQueueHeightToMountedLists(overlay, 0, applyHeight),
|
||||
false,
|
||||
"Native height (0) must not force a max-height"
|
||||
);
|
||||
assert.equal(
|
||||
reapplyQueueHeightToMountedLists(document.createElement("div"), 360, applyHeight),
|
||||
false,
|
||||
"overlays without legacy lists must be ignored"
|
||||
);
|
||||
|
||||
console.log("queue height DOM regression test passed");
|
||||
@@ -0,0 +1,108 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.stereogram import SaturnStereo3D
|
||||
|
||||
class TestSaturnStereo3D(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = SaturnStereo3D()
|
||||
self.img = torch.rand((1, 128, 128, 3), dtype=torch.float32)
|
||||
self.depth = torch.linspace(0, 1, 128).view(1, 1, 128).expand(1, 128, 128).unsqueeze(-1).repeat(1, 1, 1, 3)
|
||||
|
||||
def test_cross_eye_mode_dimensions(self):
|
||||
# Image is 128x128. With gap=10, border=5:
|
||||
# Side by side: 128 + 10 + 128 = 266 width.
|
||||
# With border=5 on each side: Width = 266 + 10 = 276. Height = 128 + 10 = 138.
|
||||
res = self.node.generate(
|
||||
self.img,
|
||||
depth_map=self.depth,
|
||||
mode="Cross-Eye",
|
||||
depth_intensity=30.0,
|
||||
gap_spacing=10,
|
||||
outer_border=5,
|
||||
alignment_dots=True
|
||||
)
|
||||
stereogram, left_eye, right_eye, depth_out = res
|
||||
self.assertEqual(stereogram.shape, (1, 138, 276, 3))
|
||||
self.assertEqual(left_eye.shape, (1, 128, 128, 3))
|
||||
self.assertEqual(right_eye.shape, (1, 128, 128, 3))
|
||||
self.assertEqual(depth_out.shape, (1, 128, 128, 3))
|
||||
|
||||
def test_parallel_mode(self):
|
||||
res = self.node.generate(
|
||||
self.img,
|
||||
depth_map=self.depth,
|
||||
mode="Parallel (Wall-Eyed)",
|
||||
depth_intensity=25.0,
|
||||
gap_spacing=0,
|
||||
outer_border=0,
|
||||
alignment_dots=False
|
||||
)
|
||||
stereogram, _, _, _ = res
|
||||
self.assertEqual(stereogram.shape, (1, 128, 256, 3))
|
||||
|
||||
def test_anaglyph_mode(self):
|
||||
res = self.node.generate(
|
||||
self.img,
|
||||
depth_map=self.depth,
|
||||
mode="Red-Cyan Anaglyph",
|
||||
depth_intensity=25.0
|
||||
)
|
||||
stereogram, left_eye, right_eye, _ = res
|
||||
self.assertEqual(stereogram.shape, (1, 128, 128, 3))
|
||||
np.testing.assert_allclose(stereogram[..., 0].cpu().numpy(), left_eye[..., 0].cpu().numpy(), atol=1e-5)
|
||||
np.testing.assert_allclose(stereogram[..., 1].cpu().numpy(), right_eye[..., 1].cpu().numpy(), atol=1e-5)
|
||||
np.testing.assert_allclose(stereogram[..., 2].cpu().numpy(), right_eye[..., 2].cpu().numpy(), atol=1e-5)
|
||||
|
||||
def test_top_bottom_mode(self):
|
||||
res = self.node.generate(
|
||||
self.img,
|
||||
depth_map=self.depth,
|
||||
mode="Top-Bottom (Over-Under)",
|
||||
gap_spacing=6,
|
||||
outer_border=0
|
||||
)
|
||||
stereogram, _, _, _ = res
|
||||
self.assertEqual(stereogram.shape, (1, 262, 128, 3))
|
||||
|
||||
def test_single_eye_modes(self):
|
||||
res_left = self.node.generate(self.img, depth_map=self.depth, mode="Left Eye Only")
|
||||
res_right = self.node.generate(self.img, depth_map=self.depth, mode="Right Eye Only")
|
||||
self.assertEqual(res_left[0].shape, (1, 128, 128, 3))
|
||||
self.assertEqual(res_right[0].shape, (1, 128, 128, 3))
|
||||
|
||||
def test_zoom_and_pan(self):
|
||||
res = self.node.generate(
|
||||
self.img,
|
||||
depth_map=self.depth,
|
||||
image_zoom=150.0,
|
||||
pan_horizontal=20.0,
|
||||
pan_vertical=-10.0
|
||||
)
|
||||
stereogram = res[0]
|
||||
self.assertIsNotNone(stereogram)
|
||||
self.assertEqual(stereogram.dim(), 4)
|
||||
|
||||
def test_auto_depth_fallback(self):
|
||||
from unittest.mock import patch
|
||||
with patch("nodes.stereogram.get_ort_session", side_effect=RuntimeError("onnxruntime unavailable")):
|
||||
res = self.node.generate(
|
||||
self.img,
|
||||
depth_map=None,
|
||||
mode="Cross-Eye",
|
||||
depth_intensity=20.0
|
||||
)
|
||||
stereogram, _, _, depth_out = res
|
||||
self.assertEqual(stereogram.shape[0], 1)
|
||||
self.assertEqual(depth_out.shape, (1, 128, 128, 3))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.text_replacer import MultiTextReplacer, parse_find_targets
|
||||
|
||||
class TestMultiTextReplacer(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = MultiTextReplacer()
|
||||
|
||||
def test_parse_find_targets_basic(self):
|
||||
targets = parse_find_targets("dog, cat, bird")
|
||||
self.assertEqual(targets, ["dog", "cat", "bird"])
|
||||
|
||||
def test_parse_find_targets_with_quotes(self):
|
||||
targets = parse_find_targets('"dog", "cat", "bird"')
|
||||
self.assertEqual(targets, ["dog", "cat", "bird"])
|
||||
|
||||
def test_parse_find_targets_with_comma_inside_quotes(self):
|
||||
targets = parse_find_targets('"hello, world", "foo, bar", baz')
|
||||
self.assertEqual(targets, ["hello, world", "foo, bar", "baz"])
|
||||
|
||||
def test_replace_comma_separated_list(self):
|
||||
text, count = self.node.replace_text(
|
||||
find="dog, cat, bird",
|
||||
replace="animal",
|
||||
case_sensitive=False,
|
||||
search_mode="Comma Separated List",
|
||||
text="I have a dog and a cat and a bird."
|
||||
)
|
||||
self.assertEqual(text, "I have a animal and a animal and a animal.")
|
||||
self.assertEqual(count, 3)
|
||||
|
||||
def test_replace_longest_match_first(self):
|
||||
text, count = self.node.replace_text(
|
||||
find="dog, hot dog",
|
||||
replace="food",
|
||||
search_mode="Comma Separated List",
|
||||
text="I ate a hot dog."
|
||||
)
|
||||
self.assertEqual(text, "I ate a food.")
|
||||
self.assertEqual(count, 1)
|
||||
|
||||
def test_case_sensitivity(self):
|
||||
text_ci, count_ci = self.node.replace_text(
|
||||
find="cat",
|
||||
replace="feline",
|
||||
case_sensitive=False,
|
||||
text="Cat and cat"
|
||||
)
|
||||
self.assertEqual(text_ci, "feline and feline")
|
||||
self.assertEqual(count_ci, 2)
|
||||
|
||||
text_cs, count_cs = self.node.replace_text(
|
||||
find="cat",
|
||||
replace="feline",
|
||||
case_sensitive=True,
|
||||
text="Cat and cat"
|
||||
)
|
||||
self.assertEqual(text_cs, "Cat and feline")
|
||||
self.assertEqual(count_cs, 1)
|
||||
|
||||
def test_regex_mode(self):
|
||||
text, count = self.node.replace_text(
|
||||
find=r"\b\d{4}\b",
|
||||
replace="YEAR",
|
||||
search_mode="Regex Pattern",
|
||||
text="In 1999 and 2024."
|
||||
)
|
||||
self.assertEqual(text, "In YEAR and YEAR.")
|
||||
self.assertEqual(count, 2)
|
||||
|
||||
def test_edge_case_empty_text(self):
|
||||
text, count = self.node.replace_text(find="test", replace="new", text="")
|
||||
self.assertEqual(text, "")
|
||||
self.assertEqual(count, 0)
|
||||
|
||||
def test_edge_case_none_text(self):
|
||||
text, count = self.node.replace_text(find="test", replace="new", text=None)
|
||||
self.assertEqual(text, "")
|
||||
self.assertEqual(count, 0)
|
||||
|
||||
def test_edge_case_empty_find(self):
|
||||
text, count = self.node.replace_text(find="", replace="new", text="Hello")
|
||||
self.assertEqual(text, "Hello")
|
||||
self.assertEqual(count, 0)
|
||||
|
||||
def test_edge_case_deletion_empty_replace(self):
|
||||
text, count = self.node.replace_text(find="badword", replace="", text="badword good text")
|
||||
self.assertEqual(text, " good text")
|
||||
self.assertEqual(count, 1)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,114 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.text_split import SaturnTextSplit
|
||||
|
||||
class TestSaturnTextSplit(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = SaturnTextSplit()
|
||||
|
||||
def test_forward_literal_split_basic(self):
|
||||
# User standard case: "Hello --- Bye" with delimiter "--"
|
||||
t1, t2 = self.node.split_text(
|
||||
text="Hello --- Bye",
|
||||
split_by="--",
|
||||
use_regex=False,
|
||||
split_direction="forward (first occurrence from start)",
|
||||
strip_whitespace=False
|
||||
)
|
||||
self.assertEqual(t1, "Hello ")
|
||||
self.assertEqual(t2, "- Bye")
|
||||
|
||||
def test_backward_literal_split_basic(self):
|
||||
# Splitting from end
|
||||
t1, t2 = self.node.split_text(
|
||||
text="Hello --- Bye",
|
||||
split_by="--",
|
||||
use_regex=False,
|
||||
split_direction="backward (last occurrence from end)",
|
||||
strip_whitespace=False
|
||||
)
|
||||
self.assertEqual(t1, "Hello -")
|
||||
self.assertEqual(t2, " Bye")
|
||||
|
||||
def test_strip_whitespace(self):
|
||||
t1, t2 = self.node.split_text(
|
||||
text=" Left Part --- Right Part ",
|
||||
split_by="---",
|
||||
use_regex=False,
|
||||
strip_whitespace=True
|
||||
)
|
||||
self.assertEqual(t1, "Left Part")
|
||||
self.assertEqual(t2, "Right Part")
|
||||
|
||||
def test_forward_regex_split(self):
|
||||
t1, t2 = self.node.split_text(
|
||||
text="prompt_v1__step_100__seed_42",
|
||||
split_by=r"__step_\d+__",
|
||||
use_regex=True,
|
||||
split_direction="forward (first occurrence from start)"
|
||||
)
|
||||
self.assertEqual(t1, "prompt_v1")
|
||||
self.assertEqual(t2, "seed_42")
|
||||
|
||||
def test_backward_regex_split(self):
|
||||
t1, t2 = self.node.split_text(
|
||||
text="tag1#123tag2#456tag3",
|
||||
split_by=r"#\d+",
|
||||
use_regex=True,
|
||||
split_direction="backward (last occurrence from end)"
|
||||
)
|
||||
self.assertEqual(t1, "tag1#123tag2")
|
||||
self.assertEqual(t2, "tag3")
|
||||
|
||||
def test_edge_case_empty_text(self):
|
||||
t1, t2 = self.node.split_text(text="", split_by="--")
|
||||
self.assertEqual(t1, "")
|
||||
self.assertEqual(t2, "")
|
||||
|
||||
def test_edge_case_none_text(self):
|
||||
t1, t2 = self.node.split_text(text=None, split_by="--")
|
||||
self.assertEqual(t1, "")
|
||||
self.assertEqual(t2, "")
|
||||
|
||||
def test_edge_case_empty_delimiter(self):
|
||||
t1, t2 = self.node.split_text(text="Some text", split_by="")
|
||||
self.assertEqual(t1, "Some text")
|
||||
self.assertEqual(t2, "")
|
||||
|
||||
def test_edge_case_delimiter_not_found(self):
|
||||
t1, t2 = self.node.split_text(text="Hello World", split_by="---")
|
||||
self.assertEqual(t1, "Hello World")
|
||||
self.assertEqual(t2, "")
|
||||
|
||||
def test_edge_case_delimiter_at_start(self):
|
||||
t1, t2 = self.node.split_text(text="--Rest of message", split_by="--")
|
||||
self.assertEqual(t1, "")
|
||||
self.assertEqual(t2, "Rest of message")
|
||||
|
||||
def test_edge_case_delimiter_at_end(self):
|
||||
t1, t2 = self.node.split_text(text="First part--", split_by="--")
|
||||
self.assertEqual(t1, "First part")
|
||||
self.assertEqual(t2, "")
|
||||
|
||||
def test_edge_case_invalid_regex_fallback(self):
|
||||
t1, t2 = self.node.split_text(
|
||||
text="Hello [unclosed bracket World",
|
||||
split_by="[unclosed bracket",
|
||||
use_regex=True
|
||||
)
|
||||
self.assertEqual(t1, "Hello ")
|
||||
self.assertEqual(t2, " World")
|
||||
|
||||
def test_multiline_text_split(self):
|
||||
text = "Line 1\nLine 2\n---\nLine 3\nLine 4"
|
||||
t1, t2 = self.node.split_text(text=text, split_by="\n---\n")
|
||||
self.assertEqual(t1, "Line 1\nLine 2")
|
||||
self.assertEqual(t2, "Line 3\nLine 4")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,255 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import test_helper
|
||||
|
||||
from nodes.utils import (
|
||||
parse_pretty_name,
|
||||
parse_pretty_name_with_version,
|
||||
sanitize_folder_path,
|
||||
format_lora_output_name,
|
||||
is_local_request,
|
||||
is_safe_path,
|
||||
sanitize_image_loader_folder,
|
||||
get_csrf_token,
|
||||
is_same_origin_or_direct,
|
||||
is_authenticated_local_request,
|
||||
get_env_setting,
|
||||
invalidate_env_settings_cache
|
||||
)
|
||||
|
||||
class TestUtils(unittest.TestCase):
|
||||
def test_user_state_is_isolated_from_live_comfyui(self):
|
||||
from nodes.utils import get_saturnnodes_user_dir
|
||||
|
||||
expected_root = os.path.normpath(test_helper.TEST_USER_ROOT)
|
||||
actual = os.path.normpath(get_saturnnodes_user_dir())
|
||||
self.assertEqual(os.path.commonpath([expected_root, actual]), expected_root)
|
||||
self.assertNotEqual(actual, os.path.normpath(os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
"..", "..", "user", "default", "SaturnNodes")))
|
||||
|
||||
def test_is_local_request(self):
|
||||
class DummyReq:
|
||||
def __init__(self, ip):
|
||||
self.remote = ip
|
||||
self.assertTrue(is_local_request(DummyReq("127.0.0.1")))
|
||||
self.assertTrue(is_local_request(DummyReq("::1")))
|
||||
self.assertTrue(is_local_request(DummyReq("localhost")))
|
||||
self.assertTrue(is_local_request(DummyReq("::ffff:127.0.0.1")))
|
||||
self.assertTrue(is_local_request(DummyReq("[::ffff:127.0.0.1]")))
|
||||
self.assertFalse(is_local_request(DummyReq("192.168.1.50")))
|
||||
self.assertFalse(is_local_request(DummyReq("10.0.0.1")))
|
||||
self.assertFalse(is_local_request(DummyReq("8.8.8.8")))
|
||||
|
||||
def test_csrf_token_generation(self):
|
||||
token = get_csrf_token()
|
||||
self.assertIsInstance(token, str)
|
||||
self.assertEqual(len(token), 64)
|
||||
# Session token is consistent across calls in same session
|
||||
self.assertEqual(get_csrf_token(), token)
|
||||
|
||||
def test_is_same_origin_or_direct(self):
|
||||
class DummyReq:
|
||||
def __init__(self, headers=None):
|
||||
self.headers = headers or {}
|
||||
|
||||
# 1. Direct or same-origin
|
||||
self.assertTrue(is_same_origin_or_direct(DummyReq()))
|
||||
self.assertTrue(is_same_origin_or_direct(DummyReq({"Origin": "http://127.0.0.1:8188"})))
|
||||
self.assertTrue(is_same_origin_or_direct(DummyReq({"Origin": "http://localhost:8188"})))
|
||||
|
||||
# 2. Block cross-site via Sec-Fetch-Site
|
||||
self.assertFalse(is_same_origin_or_direct(DummyReq({"Sec-Fetch-Site": "cross-site"})))
|
||||
|
||||
# 3. Block malicious cross-origin
|
||||
self.assertFalse(is_same_origin_or_direct(DummyReq({"Origin": "http://attacker-site.com"})))
|
||||
self.assertFalse(is_same_origin_or_direct(DummyReq({"Referer": "http://malicious.org/exploit.html"})))
|
||||
|
||||
def test_is_authenticated_local_request(self):
|
||||
class DummyReq:
|
||||
def __init__(self, ip="127.0.0.1", headers=None):
|
||||
self.remote = ip
|
||||
self.headers = headers or {}
|
||||
|
||||
valid_token = get_csrf_token()
|
||||
|
||||
# Valid local request with matching token
|
||||
valid_req = DummyReq("127.0.0.1", {
|
||||
"Origin": "http://127.0.0.1:8188",
|
||||
"X-SaturnNodes-CSRF-Token": valid_token
|
||||
})
|
||||
self.assertTrue(is_authenticated_local_request(valid_req))
|
||||
|
||||
# Missing token
|
||||
no_token_req = DummyReq("127.0.0.1", {"Origin": "http://127.0.0.1:8188"})
|
||||
self.assertFalse(is_authenticated_local_request(no_token_req))
|
||||
|
||||
# Wrong token
|
||||
wrong_token_req = DummyReq("127.0.0.1", {
|
||||
"Origin": "http://127.0.0.1:8188",
|
||||
"X-SaturnNodes-CSRF-Token": "invalid_forged_token"
|
||||
})
|
||||
self.assertFalse(is_authenticated_local_request(wrong_token_req))
|
||||
|
||||
# Remote IP attempting to forge or replay token
|
||||
remote_req = DummyReq("192.168.1.100", {
|
||||
"Origin": "http://127.0.0.1:8188",
|
||||
"X-SaturnNodes-CSRF-Token": valid_token
|
||||
})
|
||||
self.assertFalse(is_authenticated_local_request(remote_req))
|
||||
|
||||
# Cross-site drive-by request (Sec-Fetch-Site: cross-site)
|
||||
drive_by_req = DummyReq("127.0.0.1", {
|
||||
"Sec-Fetch-Site": "cross-site",
|
||||
"X-SaturnNodes-CSRF-Token": valid_token
|
||||
})
|
||||
self.assertFalse(is_authenticated_local_request(drive_by_req))
|
||||
|
||||
def test_is_safe_path(self):
|
||||
import tempfile
|
||||
base1 = tempfile.mkdtemp()
|
||||
sub = os.path.join(base1, "images", "test.png")
|
||||
escape = os.path.join(base1, "..", "passwords.txt")
|
||||
self.assertTrue(is_safe_path(sub, allowed_bases=[base1]))
|
||||
self.assertFalse(is_safe_path(escape, allowed_bases=[base1]))
|
||||
import shutil
|
||||
shutil.rmtree(base1, ignore_errors=True)
|
||||
|
||||
def test_parse_pretty_name(self):
|
||||
name = parse_pretty_name("char_cyberpunk-character_v1.safetensors")
|
||||
self.assertEqual(name, "Cyberpunk Character")
|
||||
|
||||
def test_parse_pretty_name_special_keywords(self):
|
||||
name = parse_pretty_name("tag_cyberpunk-nsfw-v2.safetensors")
|
||||
self.assertIn("NSFW", name)
|
||||
self.assertIn("V2", name)
|
||||
|
||||
def test_parse_pretty_name_with_version(self):
|
||||
name_v = parse_pretty_name_with_version("char_cyberpunk-character_v1.safetensors")
|
||||
self.assertEqual(name_v, "Cyberpunk Character V1")
|
||||
|
||||
def test_format_lora_output_name_modes(self):
|
||||
rel_path = os.path.join("characters", "cyberpunk_character.safetensors")
|
||||
|
||||
parsed = format_lora_output_name(rel_path, "Cyberpunk Character", "Parsed Name")
|
||||
self.assertEqual(parsed, "Cyberpunk Character")
|
||||
|
||||
filename = format_lora_output_name(rel_path, "Cyberpunk Character", "Filename")
|
||||
self.assertEqual(filename, "cyberpunk_character.safetensors")
|
||||
|
||||
no_ext = format_lora_output_name(rel_path, "Cyberpunk Character", "Filename without extension")
|
||||
self.assertEqual(no_ext, "cyberpunk_character")
|
||||
|
||||
custom_reg = format_lora_output_name(rel_path, "Cyberpunk Character", "Custom Regex", custom_regex=r"cyberpunk_([a-z]+)")
|
||||
self.assertEqual(custom_reg, "cyberpunk_character")
|
||||
|
||||
def test_sanitize_folder_path_wildcard_stripping(self):
|
||||
clean = sanitize_folder_path("watch/*")
|
||||
self.assertNotIn("*", clean)
|
||||
|
||||
def test_sanitize_folder_path_blocks_absolute_paths(self):
|
||||
import folder_paths
|
||||
inp = os.path.abspath(os.path.realpath(folder_paths.get_input_directory()))
|
||||
out = os.path.abspath(os.path.realpath(folder_paths.get_output_directory()))
|
||||
|
||||
# Test Windows drive letters and root paths
|
||||
for bad_path in ["C:\\Windows\\System32", "D:\\SecretData", "/etc/passwd", "\\\\server\\share"]:
|
||||
res = sanitize_folder_path(bad_path)
|
||||
res_norm = os.path.abspath(os.path.realpath(res))
|
||||
# Must be confined strictly within input or output directory
|
||||
is_confined = (os.path.commonpath([inp, res_norm]) == inp) or (os.path.commonpath([out, res_norm]) == out)
|
||||
self.assertTrue(is_confined, f"Path '{bad_path}' escaped confinement! Got: '{res_norm}'")
|
||||
|
||||
def test_sanitize_folder_path_blocks_directory_traversal(self):
|
||||
import folder_paths
|
||||
inp = os.path.abspath(os.path.realpath(folder_paths.get_input_directory()))
|
||||
out = os.path.abspath(os.path.realpath(folder_paths.get_output_directory()))
|
||||
|
||||
for bad_traversal in ["../../etc", "watch/../../secret", "..\\..\\Windows", "sub/../../../root"]:
|
||||
res = sanitize_folder_path(bad_traversal)
|
||||
res_norm = os.path.abspath(os.path.realpath(res))
|
||||
is_confined = (os.path.commonpath([inp, res_norm]) == inp) or (os.path.commonpath([out, res_norm]) == out)
|
||||
self.assertTrue(is_confined, f"Traversal '{bad_traversal}' escaped confinement! Got: '{res_norm}'")
|
||||
|
||||
def test_sanitize_folder_path_allows_safe_subfolders(self):
|
||||
import folder_paths
|
||||
inp = os.path.abspath(os.path.realpath(folder_paths.get_input_directory()))
|
||||
out = os.path.abspath(os.path.realpath(folder_paths.get_output_directory()))
|
||||
|
||||
for safe_sub in ["watch", "my_images", "subfolder/nested"]:
|
||||
res = sanitize_folder_path(safe_sub)
|
||||
res_norm = os.path.abspath(os.path.realpath(res))
|
||||
is_confined = (os.path.commonpath([inp, res_norm]) == inp) or (os.path.commonpath([out, res_norm]) == out)
|
||||
self.assertTrue(is_confined, f"Safe subfolder '{safe_sub}' failed confinement! Got: '{res_norm}'")
|
||||
|
||||
def test_is_safe_external_image_url(self):
|
||||
from nodes.lora_loader import is_safe_external_image_url
|
||||
self.assertTrue(is_safe_external_image_url("https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/test.jpg"))
|
||||
self.assertTrue(is_safe_external_image_url("https://image.tmdb.org/t/p/w500/sample.jpg"))
|
||||
self.assertFalse(is_safe_external_image_url("file:///etc/passwd"))
|
||||
self.assertFalse(is_safe_external_image_url("file:///C:/Windows/system32/cmd.exe"))
|
||||
self.assertFalse(is_safe_external_image_url("http://127.0.0.1:8188/secret"))
|
||||
self.assertFalse(is_safe_external_image_url("http://localhost:8080/admin"))
|
||||
self.assertFalse(is_safe_external_image_url("http://169.254.169.254/latest/meta-data/"))
|
||||
self.assertFalse(is_safe_external_image_url("ftp://server/image.png"))
|
||||
self.assertFalse(is_safe_external_image_url(""))
|
||||
self.assertFalse(is_safe_external_image_url(None))
|
||||
|
||||
def test_edge_cases_empty_or_none(self):
|
||||
self.assertEqual(parse_pretty_name(""), "")
|
||||
self.assertEqual(parse_pretty_name(None), "")
|
||||
self.assertEqual(parse_pretty_name("[ NONE ]"), "")
|
||||
self.assertEqual(parse_pretty_name("[ RANDOM ]"), "")
|
||||
|
||||
def test_sanitize_image_loader_folder_confinement(self):
|
||||
from nodes.utils import sanitize_image_loader_folder
|
||||
import folder_paths
|
||||
inp = os.path.abspath(os.path.realpath(folder_paths.get_input_directory()))
|
||||
out = os.path.abspath(os.path.realpath(folder_paths.get_output_directory()))
|
||||
|
||||
for bad in ["C:\\Windows", "/etc", "../../escape", "output/../../secret", "\\\\server\\share"]:
|
||||
res = sanitize_image_loader_folder(bad)
|
||||
res_norm = os.path.abspath(os.path.realpath(res))
|
||||
is_confined = (os.path.commonpath([inp, res_norm]) == inp) or (os.path.commonpath([out, res_norm]) == out)
|
||||
self.assertTrue(is_confined, f"Image folder '{bad}' escaped confinement: '{res_norm}'")
|
||||
|
||||
def test_get_env_setting_cache_and_invalidation(self):
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
env_file = os.path.join(temp_dir, ".env")
|
||||
try:
|
||||
with open(env_file, "w", encoding="utf-8") as f:
|
||||
f.write("# comment\nA=1\nB=\nMALFORMED_LINE\nC=hello world\n")
|
||||
|
||||
with patch("nodes.utils.get_saturnnodes_user_dir", return_value=temp_dir):
|
||||
invalidate_env_settings_cache()
|
||||
self.assertEqual(get_env_setting("A", "fallback"), "1")
|
||||
self.assertEqual(get_env_setting("B", "fallback"), "")
|
||||
self.assertEqual(get_env_setting("C", "fallback"), "hello world")
|
||||
self.assertEqual(get_env_setting("MISSING", "fallback"), "fallback")
|
||||
|
||||
# Cached value is returned until mtime changes.
|
||||
with open(env_file, "w", encoding="utf-8") as f:
|
||||
f.write("A=2\n")
|
||||
future = time.time() + 5
|
||||
os.utime(env_file, (future, future))
|
||||
self.assertEqual(get_env_setting("A", "fallback"), "2")
|
||||
self.assertEqual(get_env_setting("B", "fallback"), "fallback")
|
||||
|
||||
# Explicit invalidation forces a re-read even with a forced mtime.
|
||||
with open(env_file, "w", encoding="utf-8") as f:
|
||||
f.write("A=3\n")
|
||||
invalidate_env_settings_cache()
|
||||
self.assertEqual(get_env_setting("A", "fallback"), "3")
|
||||
finally:
|
||||
invalidate_env_settings_cache()
|
||||
shutil.rmtree(temp_dir, ignore_errors=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,137 @@
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import path from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import {
|
||||
createDocument,
|
||||
stripImports,
|
||||
} from "./fake_dom.mjs";
|
||||
import { resolveLiveElement, resolveLiveWidgetElement } from "../web/js/live_dom.js";
|
||||
|
||||
const here = path.dirname(fileURLToPath(import.meta.url));
|
||||
const repoRoot = path.resolve(here, "..");
|
||||
|
||||
// 1. Resolver behavior in isolation.
|
||||
{
|
||||
const document = createDocument();
|
||||
// live_dom.js is a real ES module and resolves `document` from globalThis.
|
||||
globalThis.document = document;
|
||||
const stale = document.createElement("div");
|
||||
const fresh = document.createElement("div");
|
||||
fresh.setAttribute("data-node-id", "7");
|
||||
const widget = document.createElement("div");
|
||||
widget.setAttribute("data-widget-name", "prompt_count_preview");
|
||||
fresh.appendChild(widget);
|
||||
document.body.appendChild(fresh);
|
||||
|
||||
assert.equal(
|
||||
resolveLiveWidgetElement(stale, 7, "prompt_count_preview"),
|
||||
widget,
|
||||
"detached cache must re-resolve to the remounted widget"
|
||||
);
|
||||
|
||||
const live = document.createElement("div");
|
||||
document.body.appendChild(live);
|
||||
let findCalls = 0;
|
||||
assert.equal(
|
||||
resolveLiveElement(live, () => { findCalls++; return fresh; }),
|
||||
live,
|
||||
"live cache must be used as-is"
|
||||
);
|
||||
assert.equal(findCalls, 0, "live cache must not trigger a DOM lookup");
|
||||
|
||||
let missingCalls = 0;
|
||||
assert.equal(
|
||||
resolveLiveElement(stale, () => { missingCalls++; return null; }),
|
||||
stale,
|
||||
"pre-mount writes must keep using the cached element"
|
||||
);
|
||||
assert.equal(missingCalls, 1, "a lookup is attempted once before falling back");
|
||||
}
|
||||
|
||||
// 2. Prompt counter: a remounted badge must receive the next update.
|
||||
const document = createDocument();
|
||||
globalThis.document = document;
|
||||
const extension = { name: "", nodeCreated: null };
|
||||
const app = {
|
||||
registerExtension(candidate) {
|
||||
extension.name = candidate.name;
|
||||
extension.nodeCreated = candidate.nodeCreated;
|
||||
},
|
||||
graph: { setDirtyCanvas() {} },
|
||||
};
|
||||
|
||||
let source = stripImports(
|
||||
fs.readFileSync(path.join(repoRoot, "web", "js", "prompt_counter.js"), "utf8")
|
||||
);
|
||||
const context = vm.createContext({
|
||||
console,
|
||||
document,
|
||||
app,
|
||||
isSaturnColorsEnabled: () => false,
|
||||
resolveLiveElement,
|
||||
resolveLiveWidgetElement,
|
||||
setTimeout,
|
||||
clearTimeout,
|
||||
Math,
|
||||
JSON,
|
||||
String,
|
||||
Number,
|
||||
Object,
|
||||
Array,
|
||||
Map,
|
||||
Set,
|
||||
Promise,
|
||||
Date,
|
||||
parseInt,
|
||||
isNaN,
|
||||
Infinity,
|
||||
});
|
||||
vm.runInContext(source, context, { filename: "prompt_counter.js" });
|
||||
|
||||
assert.equal(typeof extension.nodeCreated, "function", "prompt counter extension must register");
|
||||
|
||||
const node = {
|
||||
id: 7,
|
||||
comfyClass: "PromptCounter",
|
||||
type: "PromptCounter",
|
||||
size: [240, 120],
|
||||
flags: {},
|
||||
widgets: [
|
||||
{ name: "text", value: "a\n\nb", callback: null, inputEl: null },
|
||||
{ name: "separator", value: ">1 Empty Line", callback: null },
|
||||
{ name: "custom_regex", value: "", callback: null },
|
||||
],
|
||||
addDOMWidget(name, type, element) {
|
||||
const widget = { name, type, element, value: "" };
|
||||
this.widgets.push(widget);
|
||||
return widget;
|
||||
},
|
||||
};
|
||||
|
||||
await extension.nodeCreated(node);
|
||||
|
||||
const staleSpan = node._promptBadgeSpan;
|
||||
assert.ok(staleSpan, "badge span must be created");
|
||||
assert.equal(staleSpan.textContent, "0 Prompts", "initial badge label");
|
||||
|
||||
// Simulate the Vue renderer remounting the DOM widget with a fresh subtree.
|
||||
const wrapper = document.createElement("div");
|
||||
wrapper.setAttribute("data-node-id", "7");
|
||||
const freshBadge = document.createElement("div");
|
||||
freshBadge.setAttribute("data-widget-name", "prompt_count_preview");
|
||||
const freshSpan = document.createElement("span");
|
||||
freshSpan.textContent = "0 Prompts";
|
||||
freshBadge.appendChild(freshSpan);
|
||||
wrapper.appendChild(freshBadge);
|
||||
document.body.appendChild(wrapper);
|
||||
|
||||
const textWidget = node.widgets.find(w => w.name === "text");
|
||||
textWidget.callback();
|
||||
|
||||
assert.equal(node._promptBadgeSpan, freshSpan, "cache must point at the remounted span");
|
||||
assert.equal(freshSpan.textContent, "2 Prompts", "remounted badge must receive the new count");
|
||||
assert.equal(staleSpan.textContent, "0 Prompts", "detached span must not be treated as visible");
|
||||
|
||||
console.log("widget guard DOM regression test passed");
|
||||
@@ -0,0 +1,757 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
import { authenticatedFetch } from "./js/auth_helper.js";
|
||||
|
||||
// 36 Curated Accessible, High-Contrast Cycling Colors for Batch Visual Grouping
|
||||
const BATCH_COLORS = [
|
||||
"#10b981", // 1. Emerald
|
||||
"#3b82f6", // 2. Blue
|
||||
"#8b5cf6", // 3. Purple
|
||||
"#ec4899", // 4. Pink
|
||||
"#f59e0b", // 5. Amber
|
||||
"#06b6d4", // 6. Cyan
|
||||
"#ef4444", // 7. Red
|
||||
"#14b8a6", // 8. Teal
|
||||
"#6366f1", // 9. Indigo
|
||||
"#f97316", // 10. Orange
|
||||
"#84cc16", // 11. Lime
|
||||
"#d946ef", // 12. Fuchsia
|
||||
"#0ea5e9", // 13. Sky
|
||||
"#a855f7", // 14. Violet
|
||||
"#e11d48", // 15. Rose
|
||||
"#22c55e", // 16. Green
|
||||
"#eab308", // 17. Yellow
|
||||
"#64748b", // 18. Slate
|
||||
"#2dd4bf", // 19. Turquoise
|
||||
"#fb7185", // 20. Coral
|
||||
"#38bdf8", // 21. Light Blue
|
||||
"#c084fc", // 22. Lavender
|
||||
"#f43f5e", // 23. Ruby
|
||||
"#4ade80", // 24. Mint
|
||||
"#fbbf24", // 25. Gold
|
||||
"#818cf8", // 26. Periwinkle
|
||||
"#fb923c", // 27. Tangerine
|
||||
"#a3e635", // 28. Chartreuse
|
||||
"#f472b6", // 29. Hot Pink
|
||||
"#22d3ee", // 30. Aqua
|
||||
"#e879f9", // 31. Orchid
|
||||
"#a7f3d0", // 32. Seafoam
|
||||
"#fde047", // 33. Lemon
|
||||
"#fda4af", // 34. Blush
|
||||
"#93c5fd", // 35. Cornflower
|
||||
"#c4b5fd" // 36. Lilac
|
||||
];
|
||||
|
||||
const STORAGE_KEY = "saturnnodes_batch_queue_meta";
|
||||
const BATCH_COUNTER_KEY = "saturnnodes_batch_counter";
|
||||
const SETTING_ID = "SaturnNodes.🎨 Batch Queue.01_ShowBatchLines";
|
||||
const SETTING_SNAPSHOT_GUARD = "SaturnNodes.🎨 Batch Queue.02_SnapshotGuard";
|
||||
|
||||
// In-memory registry mapping prompt_id -> batch metadata
|
||||
let batchRegistry = new Map();
|
||||
let activeBatchContext = null;
|
||||
let batchCounter = 0;
|
||||
|
||||
function loadStorage() {
|
||||
try {
|
||||
const raw = localStorage.getItem(STORAGE_KEY);
|
||||
if (raw) {
|
||||
const parsed = JSON.parse(raw);
|
||||
if (Array.isArray(parsed)) {
|
||||
for (const item of parsed) {
|
||||
if (item && item.prompt_id) {
|
||||
batchRegistry.set(String(item.prompt_id), item);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
const savedCounter = localStorage.getItem(BATCH_COUNTER_KEY);
|
||||
if (savedCounter) {
|
||||
// A negative or corrupted counter would index outside BATCH_COLORS
|
||||
// and render an undefined stroke color.
|
||||
batchCounter = Math.max(0, parseInt(savedCounter, 10) || 0);
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes BatchQueue] Failed to load batch storage:", e);
|
||||
}
|
||||
}
|
||||
|
||||
function saveStorage() {
|
||||
try {
|
||||
// Prune entries to the 250 most recently queued items. Map insertion
|
||||
// order is not a reliable recency order when metadata for an existing
|
||||
// prompt is updated in place.
|
||||
const allEntries = Array.from(batchRegistry.values());
|
||||
const pruned = allEntries
|
||||
.map((entry, index) => ({ entry, index }))
|
||||
.sort((a, b) => {
|
||||
const aTime = Number(a.entry && a.entry.queuedAt) || 0;
|
||||
const bTime = Number(b.entry && b.entry.queuedAt) || 0;
|
||||
return aTime - bTime || a.index - b.index;
|
||||
})
|
||||
.slice(-250)
|
||||
.map(item => item.entry);
|
||||
localStorage.setItem(STORAGE_KEY, JSON.stringify(pruned));
|
||||
localStorage.setItem(BATCH_COUNTER_KEY, String(batchCounter));
|
||||
} catch (e) {
|
||||
// Storage full or private mode
|
||||
}
|
||||
}
|
||||
|
||||
function getNextBatch(batchCount = 1) {
|
||||
batchCounter++;
|
||||
const colorIndex = (batchCounter - 1) % BATCH_COLORS.length;
|
||||
const color = BATCH_COLORS[colorIndex];
|
||||
const batchId = `batch_${Date.now()}_${Math.random().toString(36).substr(2, 5)}`;
|
||||
|
||||
const batch = {
|
||||
batchId,
|
||||
batchNumber: batchCounter,
|
||||
batchCount: Math.max(1, batchCount),
|
||||
color,
|
||||
colorIndex,
|
||||
queuedAt: Date.now(),
|
||||
prompts: []
|
||||
};
|
||||
return batch;
|
||||
}
|
||||
|
||||
// Inject CSS for the 1D git-graph style line
|
||||
function injectBatchStyles() {
|
||||
if (typeof document === "undefined") return;
|
||||
const styleId = "saturnnodes-batch-queue-styles";
|
||||
if (document.getElementById(styleId)) return;
|
||||
|
||||
const style = document.createElement("style");
|
||||
style.id = styleId;
|
||||
style.textContent = `
|
||||
/* SaturnNodes 1D Git-Graph Batch Line Container */
|
||||
.saturnnodes-batch-line-wrap {
|
||||
position: absolute;
|
||||
left: -9px;
|
||||
top: 0;
|
||||
bottom: 0;
|
||||
width: 9px;
|
||||
pointer-events: none;
|
||||
z-index: 50;
|
||||
display: flex;
|
||||
align-items: stretch;
|
||||
justify-content: center;
|
||||
overflow: visible;
|
||||
}
|
||||
|
||||
.saturnnodes-batch-line-svg {
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
display: block;
|
||||
overflow: visible;
|
||||
}
|
||||
|
||||
/* Ensure parent card has relative positioning and allows line in left margin */
|
||||
[data-testid="job-assets-list"] [data-job-id] {
|
||||
position: relative !important;
|
||||
overflow: visible !important;
|
||||
}
|
||||
|
||||
/* Support for classic frontend v1 queue table rows */
|
||||
.comfy-table tr[data-job-id],
|
||||
.comfy-queue-item {
|
||||
position: relative !important;
|
||||
overflow: visible !important;
|
||||
}
|
||||
`;
|
||||
document.head.appendChild(style);
|
||||
}
|
||||
|
||||
// Batch metadata comes from localStorage and the server, so never trust its
|
||||
// color string as HTML/SVG markup. Only complete 3/4/6/8-digit hex colors are
|
||||
// accepted (5 and 7-digit values are not valid CSS colors).
|
||||
const HEX_COLOR_PATTERN = /^#(?:[0-9a-fA-F]{3}|[0-9a-fA-F]{4}|[0-9a-fA-F]{6}|[0-9a-fA-F]{8})$/;
|
||||
|
||||
function normalizeBatchColor(color, fallback = BATCH_COLORS[0]) {
|
||||
if (typeof color === "string") {
|
||||
const trimmed = color.trim();
|
||||
if (HEX_COLOR_PATTERN.test(trimmed)) return trimmed;
|
||||
}
|
||||
return fallback;
|
||||
}
|
||||
|
||||
function normalizeSegmentType(segmentType) {
|
||||
return ["single", "first", "middle", "last"].includes(segmentType)
|
||||
? segmentType
|
||||
: "single";
|
||||
}
|
||||
|
||||
// Build 1D SVG line path according to segment type:
|
||||
// 'single': bracket (rounded top & bottom curving right to card)
|
||||
// 'first': rounded top curving right, straight all the way down
|
||||
// 'middle': straight line top to bottom
|
||||
// 'last': straight from top, rounded bottom curving right to card
|
||||
function buildSvgPath(segmentType) {
|
||||
// Coordinate system: width=10, height=48. Curve bends toward the card (x=9)
|
||||
switch (segmentType) {
|
||||
case "single":
|
||||
return "M 9,6 Q 2,6 2,16 L 2,32 Q 2,42 9,42";
|
||||
case "first":
|
||||
return "M 9,6 Q 2,6 2,16 L 2,48";
|
||||
case "middle":
|
||||
return "M 2,0 L 2,48";
|
||||
case "last":
|
||||
return "M 2,0 L 2,32 Q 2,42 9,42";
|
||||
default:
|
||||
return "M 9,6 Q 2,6 2,16 L 2,32 Q 2,42 9,42";
|
||||
}
|
||||
}
|
||||
|
||||
function renderBatchLineOnElement(element, segmentType, color) {
|
||||
if (!element) return;
|
||||
|
||||
let wrap = element.querySelector(".saturnnodes-batch-line-wrap");
|
||||
if (!wrap) {
|
||||
wrap = document.createElement("div");
|
||||
wrap.className = "saturnnodes-batch-line-wrap";
|
||||
element.appendChild(wrap);
|
||||
}
|
||||
|
||||
const safeSegment = normalizeSegmentType(segmentType);
|
||||
const safeColor = normalizeBatchColor(color);
|
||||
|
||||
// Avoid rewriting the SVG on every status/executing refresh when the
|
||||
// visible segment and batch color have not changed.
|
||||
const signature = `${safeSegment}|${safeColor}`;
|
||||
if (wrap.dataset.saturnnodesBatchSignature === signature && wrap.innerHTML) return;
|
||||
wrap.dataset.saturnnodesBatchSignature = signature;
|
||||
|
||||
const dPath = buildSvgPath(safeSegment);
|
||||
wrap.innerHTML = `
|
||||
<svg class="saturnnodes-batch-line-svg" viewBox="0 0 10 48" preserveAspectRatio="none">
|
||||
<path d="${dPath}" stroke="${safeColor}" stroke-width="3" stroke-linecap="round" stroke-linejoin="round" fill="none" />
|
||||
</svg>
|
||||
`;
|
||||
}
|
||||
|
||||
function removeBatchLineFromElement(element) {
|
||||
if (!element) return;
|
||||
const wrap = element.querySelector(".saturnnodes-batch-line-wrap");
|
||||
if (wrap) {
|
||||
wrap.remove();
|
||||
}
|
||||
}
|
||||
|
||||
// Check if SaturnNodes Batch Queue is enabled in settings
|
||||
function isBatchQueueEnabled() {
|
||||
try {
|
||||
if (app.extensionManager?.setting?.get) {
|
||||
const val = app.extensionManager.setting.get(SETTING_ID);
|
||||
if (val !== undefined) return val !== false;
|
||||
}
|
||||
if (app.ui?.settings?.get) {
|
||||
const val = app.ui.settings.get(SETTING_ID);
|
||||
if (val !== undefined) return val !== false;
|
||||
}
|
||||
} catch (e) {}
|
||||
return true;
|
||||
}
|
||||
|
||||
function isSnapshotGuardEnabled() {
|
||||
try {
|
||||
if (app.extensionManager?.setting?.get) {
|
||||
const val = app.extensionManager.setting.get(SETTING_SNAPSHOT_GUARD);
|
||||
if (val !== undefined) return val !== false;
|
||||
}
|
||||
if (app.ui?.settings?.get) {
|
||||
const val = app.ui.settings.get(SETTING_SNAPSHOT_GUARD);
|
||||
if (val !== undefined) return val !== false;
|
||||
}
|
||||
} catch (e) {}
|
||||
return true;
|
||||
}
|
||||
|
||||
function isSeedInput(key, val, nodeInputs) {
|
||||
const k = String(key).toLowerCase();
|
||||
// 1. Direct seed properties (seed, noise_seed, sampler_seed, seed_num, random_seed, etc.)
|
||||
if (k.includes("seed")) {
|
||||
return typeof val === "number" || (typeof val === "string" && /^-?\d+$/.test(val.trim()));
|
||||
}
|
||||
// 2. Dynamic random integer/number properties (random_value, random_int, etc.)
|
||||
if (k.includes("random") && typeof val === "number") {
|
||||
return true;
|
||||
}
|
||||
// 3. ComfyUI PrimitiveNode controlling a seed widget (contains control_after_generate or control_mode)
|
||||
if (key === "value" && nodeInputs && (nodeInputs.control_after_generate !== undefined || nodeInputs.control_mode !== undefined)) {
|
||||
return typeof val === "number" || (typeof val === "string" && /^-?\d+$/.test(String(val).trim()));
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function mergeDynamicSeeds(snapPrompt, livePrompt) {
|
||||
if (!snapPrompt) return livePrompt;
|
||||
try {
|
||||
const cloned = JSON.parse(JSON.stringify(snapPrompt));
|
||||
if (livePrompt && livePrompt.output && cloned.output) {
|
||||
for (const [nodeId, liveNode] of Object.entries(livePrompt.output)) {
|
||||
const snapNode = cloned.output[nodeId];
|
||||
if (snapNode && snapNode.inputs && liveNode.inputs) {
|
||||
for (const [key, val] of Object.entries(liveNode.inputs)) {
|
||||
if (isSeedInput(key, val, liveNode.inputs)) {
|
||||
snapNode.inputs[key] = val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Also sync seed values in workflow metadata so saved PNG metadata matches generated seeds
|
||||
if (livePrompt && livePrompt.workflow && Array.isArray(livePrompt.workflow.nodes) && cloned.workflow && Array.isArray(cloned.workflow.nodes)) {
|
||||
const liveWorkflowMap = new Map();
|
||||
for (const n of livePrompt.workflow.nodes) {
|
||||
if (n && n.id !== undefined) liveWorkflowMap.set(String(n.id), n);
|
||||
}
|
||||
for (const snapNode of cloned.workflow.nodes) {
|
||||
if (!snapNode || snapNode.id === undefined) continue;
|
||||
const liveWfNode = liveWorkflowMap.get(String(snapNode.id));
|
||||
if (liveWfNode && Array.isArray(liveWfNode.widgets_values) && Array.isArray(snapNode.widgets_values)) {
|
||||
const graphNode = app.graph?.getNodeById?.(snapNode.id) || app.graph?.getNodeById?.(Number(snapNode.id));
|
||||
for (let i = 0; i < liveWfNode.widgets_values.length; i++) {
|
||||
const liveVal = liveWfNode.widgets_values[i];
|
||||
const snapVal = snapNode.widgets_values[i];
|
||||
if (typeof liveVal === "number" && typeof snapVal === "number" && liveVal !== snapVal) {
|
||||
const widgetName = graphNode?.widgets?.[i]?.name;
|
||||
if (widgetName) {
|
||||
if (isSeedInput(widgetName, liveVal)) {
|
||||
snapNode.widgets_values[i] = liveVal;
|
||||
}
|
||||
} else {
|
||||
const apiInputs = cloned.output?.[String(snapNode.id)]?.inputs;
|
||||
if (apiInputs) {
|
||||
for (const [k, v] of Object.entries(apiInputs)) {
|
||||
if (v === liveVal && isSeedInput(k, v)) {
|
||||
snapNode.widgets_values[i] = liveVal;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return cloned;
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes BatchQueue] Error merging dynamic seeds into snapshot:", e);
|
||||
return snapPrompt;
|
||||
}
|
||||
}
|
||||
|
||||
// Sync with server if PersistentQueue is available
|
||||
async function syncBatchToServer(promptId, batchInfo) {
|
||||
try {
|
||||
await authenticatedFetch("/saturnnodes/batch_queue/sync", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
prompt_id: promptId,
|
||||
batch_info: batchInfo
|
||||
})
|
||||
}).catch(() => null);
|
||||
} catch (e) {
|
||||
// Silently ignore if server endpoint is not present
|
||||
}
|
||||
}
|
||||
|
||||
async function loadBatchFromServer() {
|
||||
try {
|
||||
if (!api || !api.fetchApi) return;
|
||||
const res = await api.fetchApi("/saturnnodes/batch_queue/data", { cache: "no-store" }).catch(() => null);
|
||||
if (res && res.ok) {
|
||||
const data = await res.json();
|
||||
if (data && typeof data === "object") {
|
||||
for (const [pid, info] of Object.entries(data)) {
|
||||
if (!pid || !info || typeof info !== "object") continue;
|
||||
const key = String(pid);
|
||||
const existing = batchRegistry.get(key);
|
||||
// Never let an older server snapshot overwrite newer local
|
||||
// metadata for the same prompt.
|
||||
if (existing && (Number(existing.queuedAt) || 0) > (Number(info.queuedAt) || 0)) {
|
||||
continue;
|
||||
}
|
||||
batchRegistry.set(key, info);
|
||||
}
|
||||
saveStorage();
|
||||
}
|
||||
}
|
||||
} catch (e) {}
|
||||
}
|
||||
|
||||
// Queue cards are virtualized/reused by the Vue queue. A card can keep the
|
||||
// same DOM element when it moves from pending -> running or when the next batch
|
||||
// starts, so the old batch line must be repainted after Vue has updated its
|
||||
// data-job-id. A single delayed update can race the async /jobs refresh; poll
|
||||
// briefly after each queue event and let the MutationObserver below catch the
|
||||
// actual attribute update as well.
|
||||
let batchVisualRefreshToken = 0;
|
||||
|
||||
function scheduleBatchVisualRefresh(durationMs = 1200) {
|
||||
const token = ++batchVisualRefreshToken;
|
||||
const deadline = Date.now() + durationMs;
|
||||
|
||||
const refresh = () => {
|
||||
if (token !== batchVisualRefreshToken) return;
|
||||
updateQueueBatchVisuals();
|
||||
if (Date.now() < deadline) {
|
||||
setTimeout(refresh, 200);
|
||||
}
|
||||
};
|
||||
|
||||
refresh();
|
||||
}
|
||||
|
||||
// Refresh visual lines in DOM for both v1 and v2
|
||||
async function updateQueueBatchVisuals() {
|
||||
if (!isBatchQueueEnabled()) {
|
||||
document.querySelectorAll(".saturnnodes-batch-line-wrap").forEach(el => el.remove());
|
||||
return;
|
||||
}
|
||||
|
||||
// Find all visible queue item elements in DOM (supports v2 job cards and v1 rows)
|
||||
const elements = Array.from(document.querySelectorAll('[data-testid="job-assets-list"] [data-job-id], .comfy-table tr[data-job-id]'));
|
||||
if (!elements.length) return;
|
||||
|
||||
// Elements are strictly in visual screen order from top to bottom.
|
||||
// Dedupe by prompt id: a duplicate row would otherwise skew the batch
|
||||
// grouping counts and make the line segments disagree with the list.
|
||||
const visiblePromptIds = [];
|
||||
const elMap = new Map();
|
||||
const duplicateElements = [];
|
||||
for (const el of elements) {
|
||||
const pid = el.getAttribute("data-job-id");
|
||||
if (!pid) continue;
|
||||
if (elMap.has(pid)) {
|
||||
// A duplicate row must not keep a batch line painted by an earlier
|
||||
// render either; only the first visible occurrence is authoritative.
|
||||
duplicateElements.push(el);
|
||||
continue;
|
||||
}
|
||||
visiblePromptIds.push(pid);
|
||||
elMap.set(pid, el);
|
||||
}
|
||||
if (!visiblePromptIds.length) return;
|
||||
for (const el of duplicateElements) {
|
||||
removeBatchLineFromElement(el);
|
||||
}
|
||||
|
||||
// Group visible prompts by batch once (O(n)) instead of filtering the whole
|
||||
// visible list for every rendered card (O(n²)).
|
||||
const batchGroups = new Map();
|
||||
for (const pid of visiblePromptIds) {
|
||||
const info = batchRegistry.get(pid);
|
||||
if (!info || !info.batchId) continue;
|
||||
if (!batchGroups.has(info.batchId)) batchGroups.set(info.batchId, []);
|
||||
batchGroups.get(info.batchId).push(pid);
|
||||
}
|
||||
|
||||
// For every prompt on screen, calculate segment type based on visual screen order:
|
||||
// - Top-most card on screen for Batch A curves at the top (┌)
|
||||
// - Bottom-most card on screen for Batch A curves at the bottom (└)
|
||||
// - 1-item Batch curves at both top and bottom (()
|
||||
// - Anything in between connects with straight lines (│)
|
||||
for (const pid of visiblePromptIds) {
|
||||
const el = elMap.get(pid);
|
||||
if (!el) continue;
|
||||
|
||||
const info = batchRegistry.get(pid);
|
||||
if (!info || !info.batchId) {
|
||||
removeBatchLineFromElement(el);
|
||||
continue;
|
||||
}
|
||||
|
||||
const batchId = info.batchId;
|
||||
const color = normalizeBatchColor(info.color);
|
||||
const batchCount = Math.max(1, Number(info.batchCount) || 1);
|
||||
const rawItemIndex = Number.isInteger(info.itemIndex) ? info.itemIndex : 0;
|
||||
// Corrupt metadata must not produce a "middle" segment for an index
|
||||
// outside the batch (e.g. -1 or 99 in a 3-item batch).
|
||||
const itemIndex = Math.min(Math.max(0, rawItemIndex), batchCount - 1);
|
||||
|
||||
// All visible items of this batch in top-to-bottom visual order
|
||||
const occurrencesOnScreen = batchGroups.get(batchId) || [];
|
||||
|
||||
if (occurrencesOnScreen.length <= 1) {
|
||||
// Only 1 item of this batch currently visible on screen
|
||||
if (batchCount === 1) {
|
||||
renderBatchLineOnElement(el, "single", color);
|
||||
} else {
|
||||
// Batch originally had multiple items, but only 1 is currently visible/unexecuted
|
||||
if (itemIndex === 0) {
|
||||
renderBatchLineOnElement(el, "first", color);
|
||||
} else if (itemIndex === (batchCount - 1)) {
|
||||
renderBatchLineOnElement(el, "last", color);
|
||||
} else {
|
||||
renderBatchLineOnElement(el, "middle", color);
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const topPidOnScreen = occurrencesOnScreen[0];
|
||||
const bottomPidOnScreen = occurrencesOnScreen[occurrencesOnScreen.length - 1];
|
||||
|
||||
if (pid === topPidOnScreen) {
|
||||
// Top card on screen -> curves at top (┌)
|
||||
renderBatchLineOnElement(el, "first", color);
|
||||
} else if (pid === bottomPidOnScreen) {
|
||||
// Bottom card on screen -> curves at bottom (└)
|
||||
renderBatchLineOnElement(el, "last", color);
|
||||
} else {
|
||||
// Middle cards on screen -> straight line (│)
|
||||
renderBatchLineOnElement(el, "middle", color);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let batchSubmissionQueue = Promise.resolve();
|
||||
// Number of queue actions whose submission turn has not settled yet. Used to
|
||||
// detect that ComfyUI will serialize this prompt later instead of immediately.
|
||||
let pendingSubmissions = 0;
|
||||
|
||||
function capturePromptAtClick() {
|
||||
// Serialize the graph now, while the canvas still shows what the user
|
||||
// clicked. ComfyUI's own loop calls graphToPrompt() much later when another
|
||||
// batch is still being submitted, so without this the text and LoRA
|
||||
// selection edited in the meantime bleed into the queued prompt.
|
||||
try {
|
||||
if (app && typeof app.graphToPrompt === "function") {
|
||||
return Promise.resolve(app.graphToPrompt()).catch(e => {
|
||||
console.warn("[SaturnNodes BatchQueue] Click-time prompt snapshot failed:", e);
|
||||
return null;
|
||||
});
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes BatchQueue] Click-time prompt snapshot failed:", e);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
// Hook queuePrompt to track batch execution and guard multi-item batches
|
||||
function setupQueueHooks() {
|
||||
if (setupQueueHooks._installed) return;
|
||||
setupQueueHooks._installed = true;
|
||||
loadStorage();
|
||||
|
||||
// 1. Hook app.queuePrompt to know when a batch is initiated
|
||||
if (app && typeof app.queuePrompt === "function") {
|
||||
const origQueuePrompt = app.queuePrompt.bind(app);
|
||||
app.queuePrompt = async function(number, batchCount = 1, queueNodeIds) {
|
||||
const count = Math.max(1, parseInt(batchCount, 10) || 1);
|
||||
|
||||
// Only ComfyUI's own processing loop or this submission queue can
|
||||
// delay serialization. When idle, ComfyUI snapshots immediately, so
|
||||
// skipping the extra graphToPrompt() keeps single clicks cheap.
|
||||
const shouldCapture = isSnapshotGuardEnabled() &&
|
||||
(pendingSubmissions > 0 || (app && app.processingQueue === true));
|
||||
const clickPromptPromise = shouldCapture ? capturePromptAtClick() : null;
|
||||
|
||||
pendingSubmissions++;
|
||||
|
||||
// Chain onto batchSubmissionQueue to ensure sequential batch submissions
|
||||
const turn = batchSubmissionQueue.then(async () => {
|
||||
const batch = getNextBatch(count);
|
||||
if (clickPromptPromise) {
|
||||
batch.clickPromptPromise = clickPromptPromise;
|
||||
}
|
||||
activeBatchContext = batch;
|
||||
saveStorage();
|
||||
try {
|
||||
return await origQueuePrompt(number, batchCount, queueNodeIds);
|
||||
} finally {
|
||||
activeBatchContext = null;
|
||||
}
|
||||
});
|
||||
batchSubmissionQueue = turn.catch(() => {});
|
||||
turn.then(
|
||||
() => { pendingSubmissions = Math.max(0, pendingSubmissions - 1); },
|
||||
() => { pendingSubmissions = Math.max(0, pendingSubmissions - 1); }
|
||||
);
|
||||
return turn;
|
||||
};
|
||||
}
|
||||
|
||||
// 2. Hook api.queuePrompt to record every generated prompt_id & enforce snapshot guard
|
||||
if (api && typeof api.queuePrompt === "function") {
|
||||
const origApiQueuePrompt = api.queuePrompt.bind(api);
|
||||
api.queuePrompt = async function(number, prompt, targets) {
|
||||
let promptToSend = prompt;
|
||||
|
||||
// Reserve the item index before the request is sent; concurrent
|
||||
// submissions must not race on the post-response queue length.
|
||||
if (!activeBatchContext) {
|
||||
activeBatchContext = getNextBatch(1);
|
||||
}
|
||||
// Capture this request's context. A concurrent request may finish
|
||||
// the batch and clear the global reference while we are awaiting
|
||||
// the server.
|
||||
const batchContext = activeBatchContext;
|
||||
if (!Number.isInteger(batchContext.nextItemIndex)) {
|
||||
batchContext.nextItemIndex = 0;
|
||||
}
|
||||
if (!Number.isInteger(batchContext.inFlight)) {
|
||||
batchContext.inFlight = 0;
|
||||
}
|
||||
const itemIndex = batchContext.nextItemIndex++;
|
||||
batchContext.inFlight++;
|
||||
|
||||
// Batch Graph Snapshot Guard:
|
||||
// Freeze the prompt the user clicked instead of the live canvas.
|
||||
// Multi-item batches additionally reuse that snapshot for every
|
||||
// later item so canvas edits cannot bleed into the rest of the batch.
|
||||
if (isSnapshotGuardEnabled()) {
|
||||
let clickPrompt = null;
|
||||
if (itemIndex === 0 && batchContext.clickPromptPromise) {
|
||||
try {
|
||||
clickPrompt = await batchContext.clickPromptPromise;
|
||||
} catch (e) {
|
||||
clickPrompt = null;
|
||||
}
|
||||
batchContext.clickPromptPromise = null;
|
||||
}
|
||||
|
||||
if (itemIndex === 0) {
|
||||
if (clickPrompt && clickPrompt.output) {
|
||||
try {
|
||||
batchContext.snapshotPrompt = JSON.parse(JSON.stringify(clickPrompt));
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes BatchQueue] Failed to snapshot prompt:", e);
|
||||
}
|
||||
// Keep randomized seeds from the live prompt so the
|
||||
// generated image matches the seed the UI queued.
|
||||
promptToSend = mergeDynamicSeeds(batchContext.snapshotPrompt, prompt);
|
||||
} else if (batchContext.batchCount > 1) {
|
||||
try {
|
||||
batchContext.snapshotPrompt = JSON.parse(JSON.stringify(prompt));
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes BatchQueue] Failed to snapshot prompt:", e);
|
||||
}
|
||||
}
|
||||
} else if (batchContext.snapshotPrompt) {
|
||||
promptToSend = mergeDynamicSeeds(batchContext.snapshotPrompt, prompt);
|
||||
}
|
||||
}
|
||||
|
||||
let res;
|
||||
try {
|
||||
res = await origApiQueuePrompt(number, promptToSend, targets);
|
||||
} catch (e) {
|
||||
batchContext.inFlight--;
|
||||
if (activeBatchContext === batchContext &&
|
||||
batchContext.nextItemIndex >= batchContext.batchCount &&
|
||||
batchContext.inFlight === 0) {
|
||||
activeBatchContext = null;
|
||||
}
|
||||
throw e;
|
||||
}
|
||||
batchContext.inFlight--;
|
||||
if (res && res.prompt_id) {
|
||||
const pid = String(res.prompt_id);
|
||||
const batchInfo = {
|
||||
prompt_id: pid,
|
||||
batchId: batchContext.batchId,
|
||||
batchNumber: batchContext.batchNumber,
|
||||
batchCount: batchContext.batchCount,
|
||||
itemIndex: itemIndex,
|
||||
color: batchContext.color,
|
||||
colorIndex: batchContext.colorIndex,
|
||||
isQueueFront: number < 0,
|
||||
queuedAt: Date.now()
|
||||
};
|
||||
|
||||
batchContext.prompts.push(pid);
|
||||
batchRegistry.set(pid, batchInfo);
|
||||
saveStorage();
|
||||
syncBatchToServer(pid, batchInfo);
|
||||
|
||||
// If this batch completed its expected prompt count, clear context
|
||||
if (activeBatchContext === batchContext &&
|
||||
batchContext.nextItemIndex >= batchContext.batchCount &&
|
||||
batchContext.inFlight === 0) {
|
||||
activeBatchContext = null;
|
||||
}
|
||||
|
||||
scheduleBatchVisualRefresh(600);
|
||||
}
|
||||
return res;
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.SaturnNodes.BatchQueue",
|
||||
async setup() {
|
||||
injectBatchStyles();
|
||||
setupQueueHooks();
|
||||
loadBatchFromServer().then(() => scheduleBatchVisualRefresh(1200));
|
||||
|
||||
// Listen for queue changes
|
||||
if (api && api.addEventListener) {
|
||||
const refreshAfterQueueEvent = () => scheduleBatchVisualRefresh();
|
||||
api.addEventListener("status", refreshAfterQueueEvent);
|
||||
api.addEventListener("execution_start", refreshAfterQueueEvent);
|
||||
api.addEventListener("executing", refreshAfterQueueEvent);
|
||||
api.addEventListener("execution_success", refreshAfterQueueEvent);
|
||||
api.addEventListener("execution_error", refreshAfterQueueEvent);
|
||||
api.addEventListener("execution_interrupted", refreshAfterQueueEvent);
|
||||
}
|
||||
|
||||
// Set up MutationObserver on body to detect queue sidebar rendering.
|
||||
// `attributes` is essential: Vue reuses a running card and only updates
|
||||
// its data-job-id when the next batch item starts. Without watching that
|
||||
// attribute, the old batch line can remain painted on the reused card.
|
||||
// `addedNodes` is inspected because a legacy .comfy-table can be
|
||||
// inserted with its rows already built, producing no later mutation
|
||||
// inside the table.
|
||||
const QUEUE_MUTATION_SELECTOR =
|
||||
'[data-testid="job-assets-list"], .comfy-table, [data-job-id]';
|
||||
const mutationTouchesQueue = (m) => {
|
||||
if (m.target && m.target.nodeType === 1) {
|
||||
if (m.target.closest?.(QUEUE_MUTATION_SELECTOR) ||
|
||||
m.target.matches?.(QUEUE_MUTATION_SELECTOR)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
for (const node of m.addedNodes || []) {
|
||||
if (node.nodeType !== 1) continue;
|
||||
if (node.matches?.(QUEUE_MUTATION_SELECTOR) ||
|
||||
node.querySelector?.(QUEUE_MUTATION_SELECTOR)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
let debounceTimer = null;
|
||||
const observer = new MutationObserver((mutations) => {
|
||||
let shouldUpdate = false;
|
||||
for (const m of mutations) {
|
||||
if (mutationTouchesQueue(m)) {
|
||||
shouldUpdate = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (shouldUpdate) {
|
||||
clearTimeout(debounceTimer);
|
||||
debounceTimer = setTimeout(updateQueueBatchVisuals, 60);
|
||||
}
|
||||
});
|
||||
|
||||
observer.observe(document.body, {
|
||||
childList: true,
|
||||
subtree: true,
|
||||
attributes: true,
|
||||
attributeFilter: ["data-job-id"],
|
||||
});
|
||||
|
||||
// Initial checks after load, including the async /jobs hydration.
|
||||
scheduleBatchVisualRefresh(2500);
|
||||
}
|
||||
});
|
||||
|
||||
export { BATCH_COLORS, buildSvgPath };
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,23 +0,0 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
|
||||
const COLOR_MAP = {
|
||||
"PauseQueueNode": { color: "#059669", bgcolor: "#047857" },
|
||||
"PersistentQueueNode": { color: "#059669", bgcolor: "#047857" },
|
||||
"FolderLoraLoaderVisualPrettyV2": { color: "#059669", bgcolor: "#047857" },
|
||||
"ImageLoaderVisualPrettyV2": { color: "#059669", bgcolor: "#047857" },
|
||||
"LoadRecentOutputs": { color: "#059669", bgcolor: "#047857" },
|
||||
"AutoWatcherNode": { color: "#d97706", bgcolor: "#b45309" },
|
||||
"UndoPlaceholder": { color: "#d97706", bgcolor: "#b45309" },
|
||||
"PreviewLatentLive": { color: "#7c3aed", bgcolor: "#6d28d9" }
|
||||
};
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.FlowControl.Colors",
|
||||
async nodeCreated(node) {
|
||||
if (node && node.type && COLOR_MAP[node.type]) {
|
||||
const theme = COLOR_MAP[node.type];
|
||||
node.color = theme.color;
|
||||
node.bgcolor = theme.bgcolor;
|
||||
}
|
||||
}
|
||||
});
|
||||
@@ -1,254 +0,0 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.FlowControl.Settings",
|
||||
async setup() {
|
||||
// 1. Enable Civitai Scraping Toggle
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.EnableCivitaiScraping",
|
||||
name: "🍃 FlowControl: Enable Civitai Auto-Scraping",
|
||||
type: "boolean",
|
||||
defaultValue: true,
|
||||
tooltip: "Enable/disable downloading LoRA preview thumbnails from Civitai.",
|
||||
onChange(value) {
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ enable_civitai_scraping: value ? "true" : "false" })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// 2. Civitai API Key Setting with Password Masking & Eye Icon Toggle
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.CivitaiAPIKey",
|
||||
name: "🍃 FlowControl: Civitai API Key",
|
||||
type: "text",
|
||||
defaultValue: "",
|
||||
tooltip: "Optional. Civitai API works publicly without a key for normal models. Only required for NSFW models, private models, or higher API rate limits.",
|
||||
onChange(value) {
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ civitai_api_key: value })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// 3. Enable TMDB Scraping Toggle
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.EnableTMDBScraping",
|
||||
name: "🍃 FlowControl: Enable TMDB Auto-Scraping",
|
||||
type: "boolean",
|
||||
defaultValue: true,
|
||||
tooltip: "Enable/disable downloading celebrity preview thumbnails from TMDB.",
|
||||
onChange(value) {
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ enable_tmdb_scraping: value ? "true" : "false" })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// 4. TMDB API Key Setting with Password Masking & Eye Icon Toggle
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.TMDBAPIKey",
|
||||
name: "🍃 FlowControl: TMDB API Key",
|
||||
type: "text",
|
||||
defaultValue: "",
|
||||
tooltip: "Optional. Only needed if auto-scraping celebrity preview thumbnails from TMDB.",
|
||||
onChange(value) {
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ tmdb_api_key: value })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// 5. Default Pause Queue State (Paused vs Running)
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.DefaultPauseState",
|
||||
name: "🍃 FlowControl: Default Pause Queue State on Launch",
|
||||
type: "combo",
|
||||
options: ["Paused", "Running"],
|
||||
defaultValue: "Paused",
|
||||
tooltip: "Choose whether the queue starts in Paused state or Running state on normal ComfyUI startup.",
|
||||
onChange(value) {
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ default_pause_state: value })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// 6. Default Pause Queue Mode (Finish vs Instant)
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.DefaultPauseMode",
|
||||
name: "🍃 FlowControl: Default Pause Queue Mode on Launch",
|
||||
type: "combo",
|
||||
options: ["Pause (Finish)", "Pause (Instant)"],
|
||||
defaultValue: "Pause (Finish)",
|
||||
tooltip: "Choose default pause behavior mode on ComfyUI startup.",
|
||||
onChange(value) {
|
||||
const modeKey = value === "Pause (Instant)" ? "instantly" : "after_finish";
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ default_pause_mode: modeKey })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// 7. Persistent Queue Restored Launch State Override
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.PersistentQueueRestoredState",
|
||||
name: "🍃 FlowControl: Persistent Queue Restored Launch State",
|
||||
type: "combo",
|
||||
options: ["Match Default", "Force Paused", "Force Running"],
|
||||
defaultValue: "Match Default",
|
||||
tooltip: "Override launch state when unfinished queue items are restored on startup (e.g., keep normal launch Running, but force Restored queue startup Paused).",
|
||||
onChange(value) {
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ persistent_queue_restored_state: value })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// 8. Toolbar Button Unpaused Color
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.PauseButtonUnpausedColor",
|
||||
name: "🍃 FlowControl: Pause Button Unpaused Color",
|
||||
type: "text",
|
||||
defaultValue: "#059669",
|
||||
tooltip: "Hex color code for the Pause toolbar button when execution is unpaused/running (default: #059669).",
|
||||
onChange(value) {
|
||||
document.documentElement.style.setProperty("--pq-unpaused-color", value || "#059669");
|
||||
}
|
||||
});
|
||||
|
||||
// 9. Toolbar Button Paused Color
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.PauseButtonPausedColor",
|
||||
name: "🍃 FlowControl: Pause Button Paused Color",
|
||||
type: "text",
|
||||
defaultValue: "#ea580c",
|
||||
tooltip: "Hex color code for the Pause toolbar button when execution is paused (default: #ea580c).",
|
||||
onChange(value) {
|
||||
document.documentElement.style.setProperty("--pq-paused-color", value || "#ea580c");
|
||||
}
|
||||
});
|
||||
|
||||
// 10. Pause Queue Toolbar Toggle Setting
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.EnablePauseQueue",
|
||||
name: "🍃 FlowControl: Enable Pause Queue Toolbar Button",
|
||||
type: "boolean",
|
||||
defaultValue: true,
|
||||
tooltip: "Displays the green/orange Pause & Continue button group in top action bar.",
|
||||
onChange(value) {
|
||||
const group = document.querySelector(".pq-button-group");
|
||||
if (group) {
|
||||
group.style.display = value ? "inline-flex" : "none";
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 11. Enable Persistent Queue Setting
|
||||
app.ui.settings.addSetting({
|
||||
id: "FlowControl.EnablePersistentQueue",
|
||||
name: "🍃 FlowControl: Enable Persistent Queue (Auto-Recovery)",
|
||||
type: "boolean",
|
||||
defaultValue: true,
|
||||
tooltip: "Automatically saves unfinished queue items to disk and restores them after restart/crash.",
|
||||
onChange(value) {
|
||||
api.fetchApi("/flow_control/settings", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ enable_persistent_queue: value ? "true" : "false" })
|
||||
}).catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
// Masking password inputs and adding color pickers in ComfyUI settings dialog
|
||||
const applyKeyMasking = () => {
|
||||
const inputs = document.querySelectorAll('input');
|
||||
inputs.forEach(input => {
|
||||
const settingRow = input.closest('tr') || input.parentElement;
|
||||
if (!settingRow) return;
|
||||
const text = settingRow.textContent || "";
|
||||
|
||||
// Password masking for API keys
|
||||
if ((text.includes("Civitai API Key") || text.includes("TMDB API Key")) && !input.dataset.masked) {
|
||||
input.dataset.masked = "true";
|
||||
input.type = "password";
|
||||
input.style.paddingRight = "2.2rem";
|
||||
|
||||
const toggleBtn = document.createElement("button");
|
||||
toggleBtn.type = "button";
|
||||
toggleBtn.innerHTML = "👁️";
|
||||
toggleBtn.style.cssText = "position: absolute; right: 0.5rem; background: none; border: none; cursor: pointer; font-size: 1rem; color: #a1a1aa; padding: 2px;";
|
||||
toggleBtn.title = "Toggle Show/Hide Key";
|
||||
|
||||
toggleBtn.onclick = (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
if (input.type === "password") {
|
||||
input.type = "text";
|
||||
toggleBtn.innerHTML = "🙈";
|
||||
} else {
|
||||
input.type = "password";
|
||||
toggleBtn.innerHTML = "👁️";
|
||||
}
|
||||
};
|
||||
|
||||
const container = input.parentElement;
|
||||
if (container) {
|
||||
container.style.position = "relative";
|
||||
container.appendChild(toggleBtn);
|
||||
}
|
||||
}
|
||||
|
||||
// Color picker popup for Pause button colors
|
||||
if ((text.includes("Pause Button Unpaused Color") || text.includes("Pause Button Paused Color")) && !input.dataset.colorPickerAttached) {
|
||||
input.dataset.colorPickerAttached = "true";
|
||||
const isUnpaused = text.includes("Unpaused Color");
|
||||
|
||||
const colorPicker = document.createElement("input");
|
||||
colorPicker.type = "color";
|
||||
colorPicker.value = input.value || (isUnpaused ? "#059669" : "#ea580c");
|
||||
colorPicker.style.cssText = "width: 28px; height: 28px; border: none; border-radius: 4px; cursor: pointer; padding: 0; background: none; vertical-align: middle; margin-left: 8px;";
|
||||
|
||||
colorPicker.oninput = (e) => {
|
||||
input.value = e.target.value;
|
||||
input.dispatchEvent(new Event("change", { bubbles: true }));
|
||||
const varName = isUnpaused ? "--pq-unpaused-color" : "--pq-paused-color";
|
||||
document.documentElement.style.setProperty(varName, e.target.value);
|
||||
};
|
||||
|
||||
input.oninput = (e) => {
|
||||
colorPicker.value = e.target.value;
|
||||
const varName = isUnpaused ? "--pq-unpaused-color" : "--pq-paused-color";
|
||||
document.documentElement.style.setProperty(varName, e.target.value);
|
||||
};
|
||||
|
||||
const container = input.parentElement;
|
||||
if (container) {
|
||||
container.style.display = "flex";
|
||||
container.style.alignItems = "center";
|
||||
container.appendChild(colorPicker);
|
||||
}
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
setInterval(() => {
|
||||
applyKeyMasking();
|
||||
}, 1000);
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,49 @@
|
||||
import { api } from "/scripts/api.js";
|
||||
|
||||
let _cachedCsrfToken = null;
|
||||
let _tokenPromise = null;
|
||||
|
||||
export async function getCsrfToken() {
|
||||
if (_cachedCsrfToken) {
|
||||
return _cachedCsrfToken;
|
||||
}
|
||||
if (_tokenPromise) {
|
||||
return _tokenPromise;
|
||||
}
|
||||
_tokenPromise = (async () => {
|
||||
try {
|
||||
let resp = await api.fetchApi("/saturnnodes/auth/token");
|
||||
if (resp.ok) {
|
||||
const data = await resp.json();
|
||||
_cachedCsrfToken = data.csrf_token || "";
|
||||
return _cachedCsrfToken;
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes Auth] Could not fetch CSRF token:", e);
|
||||
} finally {
|
||||
_tokenPromise = null;
|
||||
}
|
||||
return "";
|
||||
})();
|
||||
return _tokenPromise;
|
||||
}
|
||||
|
||||
export async function authenticatedFetch(endpoint, options = {}) {
|
||||
const token = await getCsrfToken();
|
||||
const opts = { ...options };
|
||||
opts.headers = { ...(opts.headers || {}) };
|
||||
if (token) {
|
||||
opts.headers["X-SaturnNodes-CSRF-Token"] = token;
|
||||
}
|
||||
return api.fetchApi(endpoint, opts);
|
||||
}
|
||||
|
||||
// Expose on window for all SaturnNodes extensions
|
||||
if (typeof window !== "undefined") {
|
||||
const authObj = {
|
||||
getToken: getCsrfToken,
|
||||
fetch: authenticatedFetch
|
||||
};
|
||||
window._SaturnNodesAuth = authObj;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,795 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
|
||||
/**
|
||||
* Fallback clipboard copy helper that works in non-secure HTTP and headless contexts.
|
||||
*/
|
||||
async function copyToClipboard(text) {
|
||||
if (!text) return false;
|
||||
if (navigator.clipboard && window.isSecureContext) {
|
||||
try {
|
||||
await navigator.clipboard.writeText(text);
|
||||
return true;
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes CopyPrompt] Navigator clipboard failed, trying textarea fallback:", e);
|
||||
}
|
||||
}
|
||||
try {
|
||||
const textarea = document.createElement("textarea");
|
||||
textarea.value = text;
|
||||
textarea.style.position = "fixed";
|
||||
textarea.style.left = "-9999px";
|
||||
textarea.style.top = "-9999px";
|
||||
textarea.style.opacity = "0";
|
||||
document.body.appendChild(textarea);
|
||||
textarea.focus();
|
||||
textarea.select();
|
||||
const successful = document.execCommand("copy");
|
||||
document.body.removeChild(textarea);
|
||||
return successful;
|
||||
} catch (err) {
|
||||
console.error("[SaturnNodes CopyPrompt] Clipboard copy fallback failed:", err);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
function isCopyEnabled(settingKey, defaultVal = true) {
|
||||
try {
|
||||
let val;
|
||||
if (app.extensionManager?.setting?.get) {
|
||||
val = app.extensionManager.setting.get(settingKey);
|
||||
}
|
||||
if ((val === undefined || val === null || val === "") && app.ui?.settings?.getSettingValue) {
|
||||
val = app.ui.settings.getSettingValue(settingKey);
|
||||
}
|
||||
if ((val === undefined || val === null || val === "") && app.ui?.settings?.get) {
|
||||
val = app.ui.settings.get(settingKey);
|
||||
}
|
||||
if (val === undefined || val === null || val === "") {
|
||||
return defaultVal;
|
||||
}
|
||||
return val === true || val === "true" || val === 1 || val === "1";
|
||||
} catch (_) {
|
||||
return defaultVal;
|
||||
}
|
||||
}
|
||||
|
||||
const INSPECT_SVG = `<svg class="size-4 pointer-events-none" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="11" cy="11" r="8"/><line x1="21" y1="21" x2="16.65" y2="16.65"/><line x1="11" y1="8" x2="11" y2="14"/><line x1="8" y1="11" x2="14" y2="11"/></svg>`;
|
||||
|
||||
function updateButtonGroupBorders(container) {
|
||||
if (!container) return;
|
||||
const buttons = Array.from(container.children).filter(el => el.tagName === "BUTTON");
|
||||
if (buttons.length === 0) return;
|
||||
if (buttons.length === 1) {
|
||||
buttons[0].classList.remove("rounded-l-lg", "rounded-r-lg", "rounded-none", "rounded-l-none", "rounded-r-none", "border-r", "border-modal-card-badge-border");
|
||||
buttons[0].classList.add("rounded-lg");
|
||||
return;
|
||||
}
|
||||
buttons.forEach((btn, index) => {
|
||||
btn.classList.remove("rounded-lg", "rounded-l-lg", "rounded-r-lg", "rounded-none", "rounded-l-none", "rounded-r-none", "border-r", "border-modal-card-badge-border");
|
||||
if (index === 0) {
|
||||
btn.classList.add("rounded-l-lg", "rounded-r-none", "border-r", "border-modal-card-badge-border");
|
||||
} else if (index === buttons.length - 1) {
|
||||
btn.classList.add("rounded-r-lg", "rounded-l-none");
|
||||
} else {
|
||||
btn.classList.add("rounded-none", "border-r", "border-modal-card-badge-border");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Client-side binary PNG chunk parser to extract prompt/parameters directly from ArrayBuffer.
|
||||
*/
|
||||
async function extractPromptFromImageUrl(imgSrc) {
|
||||
try {
|
||||
const resp = await fetch(imgSrc);
|
||||
if (!resp.ok) return null;
|
||||
const arrayBuffer = await resp.arrayBuffer();
|
||||
const dataView = new DataView(arrayBuffer);
|
||||
|
||||
// Verify PNG magic header: 0x89504E47 0x0D0A1A0A
|
||||
if (dataView.byteLength < 16) return null;
|
||||
if (dataView.getUint32(0) !== 0x89504E47 || dataView.getUint32(4) !== 0x0D0A1A0A) {
|
||||
return null;
|
||||
}
|
||||
|
||||
let offset = 8;
|
||||
const textDecoder = new TextDecoder("utf-8");
|
||||
const metadata = {};
|
||||
|
||||
while (offset < dataView.byteLength - 8) {
|
||||
const length = dataView.getUint32(offset);
|
||||
const type = String.fromCharCode(
|
||||
dataView.getUint8(offset + 4),
|
||||
dataView.getUint8(offset + 5),
|
||||
dataView.getUint8(offset + 6),
|
||||
dataView.getUint8(offset + 7)
|
||||
);
|
||||
|
||||
if (type === "tEXt" || type === "iTXt") {
|
||||
const chunkData = new Uint8Array(arrayBuffer, offset + 8, length);
|
||||
let nullIndex = 0;
|
||||
while (nullIndex < chunkData.length && chunkData[nullIndex] !== 0) nullIndex++;
|
||||
const keyword = textDecoder.decode(chunkData.subarray(0, nullIndex));
|
||||
|
||||
let textValue = "";
|
||||
if (type === "tEXt") {
|
||||
textValue = textDecoder.decode(chunkData.subarray(nullIndex + 1));
|
||||
} else if (type === "iTXt") {
|
||||
const compFlag = chunkData[nullIndex + 1];
|
||||
let textStart = nullIndex + 3;
|
||||
while (textStart < chunkData.length && chunkData[textStart] !== 0) textStart++;
|
||||
textStart++;
|
||||
while (textStart < chunkData.length && chunkData[textStart] !== 0) textStart++;
|
||||
textStart++;
|
||||
if (compFlag === 0) {
|
||||
textValue = textDecoder.decode(chunkData.subarray(textStart));
|
||||
}
|
||||
}
|
||||
if (keyword && textValue) {
|
||||
metadata[keyword.toLowerCase()] = textValue;
|
||||
}
|
||||
}
|
||||
offset += 12 + length;
|
||||
}
|
||||
|
||||
// 1. Check A1111 / WebUI / Forge / Civitai parameters
|
||||
if (metadata["parameters"]) {
|
||||
const lines = metadata["parameters"].split("\n");
|
||||
const pos = [];
|
||||
for (const line of lines) {
|
||||
if (line.trim().startsWith("Negative prompt:") || line.trim().startsWith("Steps:")) break;
|
||||
pos.push(line);
|
||||
}
|
||||
const res = pos.join("\n").trim();
|
||||
if (res) return res;
|
||||
}
|
||||
|
||||
// 2. Check native ComfyUI prompt graph JSON
|
||||
if (metadata["prompt"]) {
|
||||
try {
|
||||
const promptData = JSON.parse(metadata["prompt"]);
|
||||
if (typeof promptData === "object" && promptData !== null) {
|
||||
const startNodes = [];
|
||||
for (const [nid, ndata] of Object.entries(promptData)) {
|
||||
const ctype = ndata?.class_type || "";
|
||||
const inputs = ndata?.inputs || {};
|
||||
if (ctype.includes("Sampler") || ctype.includes("KSampler") || ctype.includes("Guider") || ctype.includes("CFG")) {
|
||||
for (const linkKey of ["positive", "conditioning", "guider"]) {
|
||||
const link = inputs[linkKey];
|
||||
if (Array.isArray(link) && link.length > 0) {
|
||||
startNodes.push(String(link[0]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const visited = new Set();
|
||||
const texts = [];
|
||||
|
||||
function walk(nid) {
|
||||
if (!nid || visited.has(nid)) return;
|
||||
visited.add(nid);
|
||||
const ndata = promptData[nid];
|
||||
if (!ndata || typeof ndata !== "object") return;
|
||||
const inputs = ndata.inputs || {};
|
||||
|
||||
for (const tkey of ["text", "prompt", "text_g", "text_l", "text_positive", "positive_prompt", "value", "string"]) {
|
||||
const val = inputs[tkey];
|
||||
if (typeof val === "string" && val.trim()) {
|
||||
const s = val.trim();
|
||||
if (!texts.includes(s)) texts.push(s);
|
||||
} else if (Array.isArray(val) && val.length > 0) {
|
||||
walk(String(val[0]));
|
||||
}
|
||||
}
|
||||
|
||||
for (const [inName, inVal] of Object.entries(inputs)) {
|
||||
if (Array.isArray(inVal) && inVal.length > 0) {
|
||||
const low = inName.toLowerCase();
|
||||
if (low.includes("negative")) continue;
|
||||
if (["conditioning", "positive", "cond", "text", "prompt", "guider"].some(k => low.includes(k))) {
|
||||
walk(String(inVal[0]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const sn of startNodes) {
|
||||
walk(sn);
|
||||
}
|
||||
|
||||
if (texts.length) return texts.join("\n");
|
||||
|
||||
// Fallback: search all text / prompt nodes
|
||||
for (const [nid, ndata] of Object.entries(promptData)) {
|
||||
const ctype = ndata?.class_type || "";
|
||||
const inputs = ndata?.inputs || {};
|
||||
if (ctype.includes("CLIPTextEncode") || ctype.includes("Text") || ctype.includes("Prompt")) {
|
||||
for (const tkey of ["text", "prompt", "text_g", "text_l"]) {
|
||||
const t = inputs[tkey];
|
||||
if (typeof t === "string" && t.trim().length > 1) {
|
||||
const s = t.trim();
|
||||
if (!texts.includes(s)) texts.push(s);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (texts.length) return texts.join("\n");
|
||||
}
|
||||
} catch (_) {}
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes CopyPrompt] Direct PNG metadata extraction failed:", e);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper to fetch image metadata and extract the positive prompt text.
|
||||
*/
|
||||
|
||||
async function getImagePromptAndMeta(imgSrc) {
|
||||
if (!imgSrc) return null;
|
||||
let url = new URL(imgSrc, window.location.origin);
|
||||
let filename = url.searchParams.get("filename");
|
||||
let type = url.searchParams.get("type") || "output";
|
||||
let subfolder = url.searchParams.get("subfolder") || "";
|
||||
|
||||
if (!filename) {
|
||||
const parts = url.pathname.split("/");
|
||||
filename = decodeURIComponent(parts[parts.length - 1]);
|
||||
}
|
||||
|
||||
let promptText = null;
|
||||
try {
|
||||
promptText = await extractPromptFromImageUrl(imgSrc);
|
||||
} catch (_) {}
|
||||
|
||||
if (!promptText && filename) {
|
||||
try {
|
||||
const query = `filename=${encodeURIComponent(filename)}&type=${encodeURIComponent(type)}&subfolder=${encodeURIComponent(subfolder)}`;
|
||||
const response = await api.fetchApi(`/saturnnodes/get_image_prompt?${query}`);
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
if (data && data.prompt) promptText = data.prompt;
|
||||
}
|
||||
} catch (_) {}
|
||||
}
|
||||
|
||||
return {
|
||||
promptText: promptText || "",
|
||||
filename: filename || "image",
|
||||
subfolder: subfolder || "",
|
||||
type: type || "output"
|
||||
};
|
||||
}
|
||||
|
||||
async function copyImagePrompt(imgSrc) {
|
||||
if (!imgSrc) return false;
|
||||
let promptText = null;
|
||||
|
||||
// 1. Try client-side direct binary extraction first
|
||||
try {
|
||||
promptText = await extractPromptFromImageUrl(imgSrc);
|
||||
} catch (_) {}
|
||||
|
||||
// 2. Fallback to server endpoint
|
||||
if (!promptText) {
|
||||
try {
|
||||
let url = new URL(imgSrc, window.location.origin);
|
||||
let filename = url.searchParams.get("filename");
|
||||
let type = url.searchParams.get("type") || "output";
|
||||
let subfolder = url.searchParams.get("subfolder") || "";
|
||||
|
||||
if (!filename) {
|
||||
const parts = url.pathname.split("/");
|
||||
filename = decodeURIComponent(parts[parts.length - 1]);
|
||||
}
|
||||
|
||||
if (filename) {
|
||||
const query = `filename=${encodeURIComponent(filename)}&type=${encodeURIComponent(type)}&subfolder=${encodeURIComponent(subfolder)}`;
|
||||
const response = await api.fetchApi(`/saturnnodes/get_image_prompt?${query}`);
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
if (data && data.prompt) {
|
||||
promptText = data.prompt;
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes] Direct prompt fetch failed:", e);
|
||||
}
|
||||
}
|
||||
|
||||
if (promptText) {
|
||||
const copied = await copyToClipboard(promptText);
|
||||
if (copied && app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "success",
|
||||
summary: "📋 Prompt Copied",
|
||||
detail: promptText.length > 80 ? promptText.slice(0, 80) + "..." : promptText,
|
||||
life: 3000
|
||||
});
|
||||
}
|
||||
return copied;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// 1. Hook into Node Context Menu (Right Click)
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.SaturnNodes.CopyPrompt",
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
const origGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions;
|
||||
nodeType.prototype.getExtraMenuOptions = function (_, options) {
|
||||
if (origGetExtraMenuOptions) origGetExtraMenuOptions.apply(this, arguments);
|
||||
|
||||
if (!isCopyEnabled("SaturnNodes.📋 Prompt Actions.02_EnableContextMenuCopyPrompt")) return;
|
||||
|
||||
const imgs = this.imgs || this.images || (this.widgets && this.widgets.filter(w => w.type === "image" || w.name === "image"));
|
||||
if (imgs && imgs.length > 0) {
|
||||
options.push({
|
||||
content: "📋 Copy Prompt",
|
||||
callback: async () => {
|
||||
const img = imgs[this.imageIndex || 0];
|
||||
const src = typeof img === "string" ? img : (img?.src || img?.value);
|
||||
if (!src) return;
|
||||
const success = await copyImagePrompt(src);
|
||||
if (!success) {
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "warn",
|
||||
summary: "SaturnNodes Copy Prompt",
|
||||
detail: "Could not extract positive prompt from this image.",
|
||||
life: 3000
|
||||
});
|
||||
} else {
|
||||
alert("SaturnNodes: Could not extract positive prompt from this image.");
|
||||
}
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
};
|
||||
}
|
||||
});
|
||||
|
||||
// 2. Active Interaction Tracking & Image URL Resolution
|
||||
let lastInteractedAssetCard = null;
|
||||
let lastInteractedAssetKey = null;
|
||||
|
||||
const ASSET_CARD_SELECTOR =
|
||||
"div[data-virtual-grid-item], [data-asset-id], .asset-card, [data-testid='asset-card'], [data-node-id], .lg-node, .comfy-image-preview, .group";
|
||||
|
||||
function assetIdentityKey(card) {
|
||||
if (!card) return null;
|
||||
const assetId = card.getAttribute?.("data-asset-id")
|
||||
|| card.closest?.("[data-asset-id]")?.getAttribute?.("data-asset-id");
|
||||
if (assetId) return `asset:${assetId}`;
|
||||
const img = card.querySelector?.("img");
|
||||
if (img && img.src) return `src:${img.src}`;
|
||||
return null;
|
||||
}
|
||||
|
||||
function isQueueItemElement(el) {
|
||||
if (!el || typeof el.closest !== "function") return false;
|
||||
// An asset card is explicitly an asset card in the gallery/history, never a queue job row
|
||||
if (el.closest("div[data-virtual-grid-item], [data-asset-id], .asset-card, [data-testid='asset-card']")) {
|
||||
return !!el.closest("[data-job-id]");
|
||||
}
|
||||
return !!el.closest("[data-job-id], .comfy-queue-item, .queue-item, .queue-list-item, .queue-entry");
|
||||
}
|
||||
|
||||
function registerAssetInteraction(target) {
|
||||
if (!target) return;
|
||||
const card = target.closest?.(ASSET_CARD_SELECTOR);
|
||||
if (card && !isQueueItemElement(card)) {
|
||||
lastInteractedAssetCard = card;
|
||||
lastInteractedAssetKey = assetIdentityKey(card);
|
||||
}
|
||||
}
|
||||
|
||||
// The virtual grid recycles card elements. A stored card can therefore be
|
||||
// detached, or now display a different asset than the one the user clicked.
|
||||
function getLiveInteractedCard() {
|
||||
const card = lastInteractedAssetCard;
|
||||
if (!card) return null;
|
||||
if (card.isConnected === false) {
|
||||
lastInteractedAssetCard = null;
|
||||
lastInteractedAssetKey = null;
|
||||
return null;
|
||||
}
|
||||
const key = assetIdentityKey(card);
|
||||
if (lastInteractedAssetKey && key && key !== lastInteractedAssetKey) {
|
||||
lastInteractedAssetCard = null;
|
||||
lastInteractedAssetKey = null;
|
||||
return null;
|
||||
}
|
||||
return card;
|
||||
}
|
||||
|
||||
document.addEventListener("pointerdown", (e) => registerAssetInteraction(e.target), true);
|
||||
document.addEventListener("click", (e) => registerAssetInteraction(e.target), true);
|
||||
document.addEventListener("contextmenu", (e) => registerAssetInteraction(e.target), true);
|
||||
|
||||
function getActiveImageSrc() {
|
||||
// 1. From last clicked / right-clicked / interacted asset card
|
||||
const interactedCard = getLiveInteractedCard();
|
||||
if (interactedCard) {
|
||||
const img = interactedCard.querySelector("img");
|
||||
if (img && img.src) return img.src;
|
||||
}
|
||||
// 2. From currently selected asset card in DOM
|
||||
const selected = document.querySelector(
|
||||
'div[data-virtual-grid-item] [data-selected="true"], [data-asset-id][data-selected="true"], .group[data-selected="true"]'
|
||||
);
|
||||
if (selected) {
|
||||
const img = selected.querySelector("img");
|
||||
if (img && img.src) return img.src;
|
||||
}
|
||||
// 3. Any active hover card
|
||||
const hovered = document.querySelector("div[data-virtual-grid-item]:hover, .asset-card:hover, .group:hover");
|
||||
if (hovered) {
|
||||
const img = hovered.querySelector("img");
|
||||
if (img && img.src) return img.src;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
// 3. Inject Copy Prompt & Inspect Directly Next to Download Button on Asset Cards
|
||||
function injectCopyPromptNextToDownload(downloadBtn) {
|
||||
if (!downloadBtn || !downloadBtn.parentElement) return;
|
||||
|
||||
if (isQueueItemElement(downloadBtn)) return;
|
||||
|
||||
const card = downloadBtn.closest(
|
||||
"div[data-virtual-grid-item], [data-asset-id], .asset-card, [data-testid='asset-card'], [data-node-id], .lg-node, .comfy-image-preview, .group"
|
||||
);
|
||||
if (!card) return;
|
||||
if (isQueueItemElement(card)) return;
|
||||
|
||||
const img = card.querySelector("img");
|
||||
if (!img || !img.src) return;
|
||||
if (img.classList.contains("size-8") || img.closest(".size-8, .size-10, .h-12")) return;
|
||||
|
||||
const parent = downloadBtn.parentElement;
|
||||
const enableCopy = isCopyEnabled("SaturnNodes.📋 Prompt Actions.01_EnableAssetsCopyPromptButton");
|
||||
const enableInspect = isCopyEnabled("SaturnNodes.📋 Prompt Actions.04_EnableInspectAssetButton", false);
|
||||
|
||||
if (!enableCopy && !enableInspect) return;
|
||||
|
||||
let baseClasses = downloadBtn.className
|
||||
.replace(/\brounded-[a-z0-9-]+\b/g, "")
|
||||
.replace(/\brounded\b/g, "")
|
||||
.replace(/\bborder-r\b/g, "")
|
||||
.replace(/\bborder-modal-card-badge-border\b/g, "")
|
||||
.trim();
|
||||
|
||||
if (!baseClasses || baseClasses.length < 5) {
|
||||
baseClasses = "inline-flex items-center justify-center font-medium font-inter transition-colors focus-visible:outline-hidden disabled:pointer-events-none disabled:opacity-50 border border-transparent shadow-xs cursor-pointer bg-modal-card-badge-background text-modal-card-badge-foreground hover:bg-modal-card-badge-background-hover size-8 p-0";
|
||||
}
|
||||
|
||||
const hasModernIcons = !!(
|
||||
downloadBtn.querySelector("[class*='icon-']") ||
|
||||
document.querySelector("[class*='icon-[lucide']")
|
||||
);
|
||||
|
||||
let lastBtn = downloadBtn;
|
||||
|
||||
// 1. Copy Prompt Button
|
||||
let copyBtn = parent.querySelector(".saturnnodes-hover-copy");
|
||||
if (enableCopy) {
|
||||
if (!copyBtn) {
|
||||
copyBtn = document.createElement("button");
|
||||
copyBtn.type = "button";
|
||||
copyBtn.title = "Copy Prompt";
|
||||
copyBtn.setAttribute("aria-label", "Copy Prompt");
|
||||
copyBtn.className = `saturnnodes-hover-btn saturnnodes-hover-copy ${baseClasses} shrink-0`;
|
||||
|
||||
if (hasModernIcons) {
|
||||
copyBtn.innerHTML = `<i class="icon-[lucide--copy] size-4 pointer-events-none"></i>`;
|
||||
} else {
|
||||
copyBtn.innerHTML = `<svg class="size-4 pointer-events-none" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>`;
|
||||
}
|
||||
|
||||
copyBtn.onclick = async (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
// Resolve the card/image at click time: the grid may have
|
||||
// recycled this card for another asset since injection.
|
||||
const liveCard = copyBtn.closest(ASSET_CARD_SELECTOR) || card;
|
||||
const activeImg = liveCard.querySelector("img") || img;
|
||||
if (!activeImg || !activeImg.src) return;
|
||||
const success = await copyImagePrompt(activeImg.src);
|
||||
if (hasModernIcons) {
|
||||
copyBtn.innerHTML = success
|
||||
? `<i class="icon-[lucide--check] size-4 text-emerald-600 pointer-events-none"></i>`
|
||||
: `<i class="icon-[lucide--x] size-4 text-rose-600 pointer-events-none"></i>`;
|
||||
setTimeout(() => {
|
||||
copyBtn.innerHTML = `<i class="icon-[lucide--copy] size-4 pointer-events-none"></i>`;
|
||||
}, 2000);
|
||||
} else {
|
||||
copyBtn.innerHTML = success
|
||||
? `<svg class="size-4 text-emerald-600 pointer-events-none" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><polyline points="20 6 9 17 4 12"/></svg>`
|
||||
: `<svg class="size-4 text-rose-600 pointer-events-none" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><line x1="18" y1="6" x2="6" y2="18"/><line x1="6" y1="6" x2="18" y2="18"/></svg>`;
|
||||
setTimeout(() => {
|
||||
copyBtn.innerHTML = `<svg class="size-4 pointer-events-none" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>`;
|
||||
}, 2000);
|
||||
}
|
||||
};
|
||||
parent.insertBefore(copyBtn, lastBtn.nextSibling);
|
||||
}
|
||||
lastBtn = copyBtn;
|
||||
} else if (copyBtn) {
|
||||
copyBtn.remove();
|
||||
}
|
||||
|
||||
// 2. Inspect Asset (Zoom) Button
|
||||
let inspectBtn = parent.querySelector(".saturnnodes-hover-inspect");
|
||||
if (enableInspect) {
|
||||
if (!inspectBtn) {
|
||||
inspectBtn = document.createElement("button");
|
||||
inspectBtn.type = "button";
|
||||
inspectBtn.title = "Inspect asset";
|
||||
inspectBtn.setAttribute("aria-label", "Inspect asset");
|
||||
inspectBtn.className = `saturnnodes-hover-btn saturnnodes-hover-inspect ${baseClasses} shrink-0`;
|
||||
inspectBtn.innerHTML = INSPECT_SVG;
|
||||
|
||||
inspectBtn.onclick = (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
const liveCard = inspectBtn.closest(ASSET_CARD_SELECTOR) || card;
|
||||
const target = liveCard.querySelector(".aspect-square, [class*='aspect-square'], img") || img;
|
||||
if (target) {
|
||||
target.dispatchEvent(new MouseEvent("dblclick", { bubbles: true, cancelable: true, view: window }));
|
||||
}
|
||||
};
|
||||
parent.insertBefore(inspectBtn, lastBtn.nextSibling);
|
||||
}
|
||||
lastBtn = inspectBtn;
|
||||
} else if (inspectBtn) {
|
||||
inspectBtn.remove();
|
||||
}
|
||||
|
||||
updateButtonGroupBorders(parent);
|
||||
}
|
||||
|
||||
// 4. Inject Copy Prompt into PrimeVue Context Menu Next to Download Item
|
||||
function injectContextMenuCopy(contextMenu) {
|
||||
if (!contextMenu) return;
|
||||
if (contextMenu.querySelector(".saturnnodes-contextmenu-copy")) return;
|
||||
if (!isCopyEnabled("SaturnNodes.📋 Prompt Actions.02_EnableContextMenuCopyPrompt")) return;
|
||||
|
||||
// Guard: strictly ignore top menu, menubar, command menus, and main application menus!
|
||||
if (
|
||||
contextMenu.closest?.(".comfy-command-menu, .comfy-menu, .p-menubar, [data-pc-name='menubar'], [data-pc-name='tieredmenu']") ||
|
||||
contextMenu.classList?.contains("comfy-command-menu") ||
|
||||
contextMenu.classList?.contains("p-tieredmenu") ||
|
||||
contextMenu.querySelector?.(".p-menubar-root-list, .p-tieredmenu-root-list, li[aria-label='File'], li[aria-label='New'], li[aria-label='Edit'], li[aria-label='View']")
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let downloadLi = null;
|
||||
const candidates = contextMenu.querySelectorAll('li[role="menuitem"], [data-pc-section="item"]');
|
||||
for (const item of candidates) {
|
||||
const label = (item.getAttribute("aria-label") || "").trim().toLowerCase();
|
||||
const text = (item.textContent || "").trim().toLowerCase();
|
||||
|
||||
// Strictly ignore workflow export / save actions
|
||||
if (label.includes("export") || text.includes("export") || label.includes("save as") || text.includes("save as")) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (label === "download" || label === "download image" || text === "download" || text === "download image") {
|
||||
downloadLi = item;
|
||||
break;
|
||||
}
|
||||
|
||||
if (!downloadLi && (label.includes("download") || text.includes("download"))) {
|
||||
downloadLi = item;
|
||||
}
|
||||
}
|
||||
if (!downloadLi || !downloadLi.parentElement) return;
|
||||
|
||||
const refBtn = downloadLi.querySelector("button");
|
||||
const btnClass = refBtn
|
||||
? refBtn.className
|
||||
: "relative inline-flex items-center gap-2 cursor-pointer touch-manipulation whitespace-nowrap appearance-none border-none font-medium font-inter transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:pointer-events-none disabled:opacity-50 text-secondary-foreground bg-secondary-background hover:bg-secondary-background-hover h-8 rounded-lg p-2 text-xs w-full justify-start p-contextmenu-item-link";
|
||||
|
||||
const hasModernIcons = !!(
|
||||
downloadLi.querySelector("[class*='icon-']") ||
|
||||
document.querySelector("[class*='icon-[lucide']")
|
||||
);
|
||||
|
||||
const copyLi = document.createElement("li");
|
||||
copyLi.className = "p-contextmenu-item saturnnodes-contextmenu-copy";
|
||||
copyLi.setAttribute("role", "menuitem");
|
||||
copyLi.setAttribute("aria-label", "Copy prompt");
|
||||
copyLi.setAttribute("data-pc-section", "item");
|
||||
copyLi.setAttribute("data-p-active", "false");
|
||||
copyLi.setAttribute("data-p-focused", "false");
|
||||
|
||||
const copyIconHtml = hasModernIcons
|
||||
? `<i class="icon-[lucide--copy] size-4"></i>`
|
||||
: `<svg class="size-4" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>`;
|
||||
|
||||
copyLi.innerHTML = `
|
||||
<div class="p-contextmenu-item-content" data-pc-section="itemcontent">
|
||||
<button class="${btnClass}" tabindex="-1" data-pc-section="itemlink">
|
||||
${copyIconHtml}
|
||||
<span>Copy prompt</span>
|
||||
</button>
|
||||
</div>
|
||||
`;
|
||||
|
||||
copyLi.addEventListener("click", async (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
const span = copyLi.querySelector("span");
|
||||
const src = getActiveImageSrc();
|
||||
if (src) {
|
||||
const success = await copyImagePrompt(src);
|
||||
if (span) span.textContent = success ? "Copied! ✅" : "Failed to copy ❌";
|
||||
} else {
|
||||
if (span) span.textContent = "No image found ❌";
|
||||
}
|
||||
setTimeout(() => {
|
||||
contextMenu.style.display = "none";
|
||||
document.dispatchEvent(new KeyboardEvent("keydown", { key: "Escape", bubbles: true }));
|
||||
}, 500);
|
||||
});
|
||||
|
||||
downloadLi.parentElement.insertBefore(copyLi, downloadLi.nextSibling);
|
||||
}
|
||||
|
||||
// 5. Fallback Hook for Older Frontend Action Bars
|
||||
function injectHoverCopyAction(overlayBar) {
|
||||
if (!overlayBar) return;
|
||||
if (!isCopyEnabled("SaturnNodes.📋 Prompt Actions.01_EnableAssetsCopyPromptButton")) return;
|
||||
|
||||
if (isQueueItemElement(overlayBar)) return;
|
||||
|
||||
const downloadBtn = overlayBar.querySelector('button[aria-label="Download"], button[aria-label*="ownload" i]');
|
||||
if (downloadBtn) {
|
||||
injectCopyPromptNextToDownload(downloadBtn);
|
||||
return;
|
||||
}
|
||||
|
||||
const iconGroup =
|
||||
overlayBar.querySelector(".flex.shrink-0") ||
|
||||
overlayBar.querySelector('button[aria-label="Zoom in"]')?.parentElement ||
|
||||
overlayBar.querySelector('button[aria-label*="zoom" i]')?.parentElement ||
|
||||
overlayBar;
|
||||
|
||||
const parentCard = overlayBar.closest(ASSET_CARD_SELECTOR);
|
||||
if (!parentCard) return;
|
||||
if (isQueueItemElement(parentCard)) return;
|
||||
|
||||
const img = parentCard.querySelector("img");
|
||||
if (!img || !img.src) return;
|
||||
if (img.classList.contains("size-8") || img.closest(".size-8, .size-10, .h-12")) return;
|
||||
|
||||
// Older action bars keep injected children when they are reused for the
|
||||
// next asset. Drop a button that is bound to a previous asset instead of
|
||||
// leaving it to copy the old image.
|
||||
const assetKey = assetIdentityKey(parentCard) || "";
|
||||
const existingBtn = overlayBar.querySelector(".saturnnodes-hover-copy");
|
||||
if (existingBtn && existingBtn.dataset.saturnnodesAssetKey === assetKey) return;
|
||||
if (existingBtn) existingBtn.remove();
|
||||
|
||||
const copyBtn = document.createElement("button");
|
||||
copyBtn.className =
|
||||
"saturnnodes-hover-btn saturnnodes-hover-copy relative inline-flex items-center justify-center cursor-pointer touch-manipulation appearance-none border-none text-xs font-medium font-inter transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:pointer-events-none disabled:opacity-50 bg-white text-gray-700 hover:bg-gray-100 size-8 p-0 rounded-none pointer-events-auto border-r border-gray-200 shrink-0";
|
||||
copyBtn.title = "Copy Prompt";
|
||||
copyBtn.setAttribute("aria-label", "Copy Prompt");
|
||||
copyBtn.setAttribute("data-pd-tooltip", "true");
|
||||
copyBtn.dataset.saturnnodesAssetKey = assetKey;
|
||||
copyBtn.innerHTML = "<span class='text-sm pointer-events-none'>📋</span>";
|
||||
|
||||
copyBtn.onclick = async (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
// Re-resolve at click time: the overlay may now belong to another asset.
|
||||
const liveCard = copyBtn.closest(ASSET_CARD_SELECTOR) || parentCard;
|
||||
const liveImg = liveCard.querySelector("img") || img;
|
||||
if (!liveImg || !liveImg.src) return;
|
||||
const success = await copyImagePrompt(liveImg.src);
|
||||
if (success) {
|
||||
copyBtn.innerHTML = "<span class='text-sm pointer-events-none'>✅</span>";
|
||||
setTimeout(() => (copyBtn.innerHTML = "<span class='text-sm pointer-events-none'>📋</span>"), 2000);
|
||||
} else {
|
||||
copyBtn.innerHTML = "<span class='text-sm pointer-events-none'>❌</span>";
|
||||
setTimeout(() => (copyBtn.innerHTML = "<span class='text-sm pointer-events-none'>📋</span>"), 2000);
|
||||
}
|
||||
};
|
||||
|
||||
const moreBtn =
|
||||
iconGroup.querySelector('button[aria-label="More options"]') ||
|
||||
iconGroup.querySelector('button[aria-label="More"]') ||
|
||||
iconGroup.lastElementChild;
|
||||
|
||||
if (moreBtn && moreBtn.parentElement === iconGroup) {
|
||||
iconGroup.insertBefore(copyBtn, moreBtn);
|
||||
} else {
|
||||
iconGroup.appendChild(copyBtn);
|
||||
}
|
||||
}
|
||||
|
||||
// 6. Active Scanner for Dynamic Virtual Grid & Context Menus
|
||||
function scanAndInject() {
|
||||
const downloadBtns = document.querySelectorAll(
|
||||
'button[aria-label="Download"], button[aria-label*="ownload" i]'
|
||||
);
|
||||
downloadBtns.forEach(injectCopyPromptNextToDownload);
|
||||
|
||||
const menus = document.querySelectorAll('.p-contextmenu, [data-pc-name="contextmenu"]');
|
||||
menus.forEach(injectContextMenuCopy);
|
||||
|
||||
const olderOverlays = document.querySelectorAll(
|
||||
'[data-testid="asset-card-actions"], .asset-card-overlay, .asset-item-overlay'
|
||||
);
|
||||
olderOverlays.forEach(injectHoverCopyAction);
|
||||
}
|
||||
|
||||
// Delegated hover listener for instant injection during virtual scrolling
|
||||
document.addEventListener("pointerover", (e) => {
|
||||
const btn = e.target.closest?.('button[aria-label="Download"], button[aria-label*="ownload" i]');
|
||||
if (btn) {
|
||||
injectCopyPromptNextToDownload(btn);
|
||||
return;
|
||||
}
|
||||
const card = e.target.closest?.('div[data-virtual-grid-item], [data-asset-id], .asset-card, .group');
|
||||
if (card) {
|
||||
registerAssetInteraction(card);
|
||||
const cardDl = card.querySelector('button[aria-label="Download"], button[aria-label*="ownload" i]');
|
||||
if (cardDl) injectCopyPromptNextToDownload(cardDl);
|
||||
}
|
||||
const menu = e.target.closest?.('.p-contextmenu, [data-pc-name="contextmenu"]');
|
||||
if (menu) {
|
||||
injectContextMenuCopy(menu);
|
||||
}
|
||||
}, { passive: true });
|
||||
|
||||
// MutationObserver for DOM changes
|
||||
const observer = new MutationObserver((mutations) => {
|
||||
for (const mutation of mutations) {
|
||||
if (mutation.type === "childList") {
|
||||
for (const node of mutation.addedNodes) {
|
||||
if (node.nodeType === Node.ELEMENT_NODE) {
|
||||
if (isQueueItemElement(node)) continue;
|
||||
|
||||
if (node.matches?.('button[aria-label="Download"], button[aria-label*="ownload" i]')) {
|
||||
injectCopyPromptNextToDownload(node);
|
||||
} else if (node.matches?.('.p-contextmenu, [data-pc-name="contextmenu"]')) {
|
||||
injectContextMenuCopy(node);
|
||||
} else if (node.matches?.('[data-testid="asset-card-actions"], .asset-card-overlay, .asset-item-overlay')) {
|
||||
injectHoverCopyAction(node);
|
||||
} else if (node.querySelectorAll) {
|
||||
const dlBtns = node.querySelectorAll('button[aria-label="Download"], button[aria-label*="ownload" i]');
|
||||
dlBtns.forEach(injectCopyPromptNextToDownload);
|
||||
|
||||
const ctxMenus = node.querySelectorAll('.p-contextmenu, [data-pc-name="contextmenu"]');
|
||||
ctxMenus.forEach(injectContextMenuCopy);
|
||||
|
||||
const overlays = node.querySelectorAll('[data-testid="asset-card-actions"], .asset-card-overlay, .asset-item-overlay');
|
||||
overlays.forEach(injectHoverCopyAction);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
|
||||
// Run scanner immediately, on DOM ready, and periodic sweep
|
||||
if (document.readyState === "loading") {
|
||||
document.addEventListener("DOMContentLoaded", scanAndInject);
|
||||
} else {
|
||||
scanAndInject();
|
||||
}
|
||||
setInterval(scanAndInject, 500);
|
||||
|
||||
|
||||
+63
-23
@@ -1,10 +1,11 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
import { authenticatedFetch } from "./auth_helper.js";
|
||||
|
||||
// Create global styles for the decision buttons
|
||||
const styles = document.createElement("style");
|
||||
styles.textContent = `
|
||||
.flowcontrol-decision-btn {
|
||||
.saturnnodes-decision-btn {
|
||||
margin: 4px;
|
||||
padding: 6px 12px;
|
||||
border: none;
|
||||
@@ -14,19 +15,19 @@ styles.textContent = `
|
||||
color: white;
|
||||
transition: all 0.2s;
|
||||
}
|
||||
.flowcontrol-decision-btn:disabled {
|
||||
.saturnnodes-decision-btn:disabled {
|
||||
opacity: 0.3;
|
||||
cursor: not-allowed;
|
||||
filter: grayscale(100%);
|
||||
}
|
||||
.flowcontrol-btn-continue { background-color: #2e7d32; }
|
||||
.flowcontrol-btn-continue:hover:not(:disabled) { background-color: #388e3c; }
|
||||
.flowcontrol-btn-cancel { background-color: #f57c00; }
|
||||
.flowcontrol-btn-cancel:hover:not(:disabled) { background-color: #fb8c00; }
|
||||
.flowcontrol-btn-stop { background-color: #d32f2f; }
|
||||
.flowcontrol-btn-stop:hover:not(:disabled) { background-color: #f44336; }
|
||||
.saturnnodes-btn-continue { background-color: #2e7d32; }
|
||||
.saturnnodes-btn-continue:hover:not(:disabled) { background-color: #388e3c; }
|
||||
.saturnnodes-btn-cancel { background-color: #f57c00; }
|
||||
.saturnnodes-btn-cancel:hover:not(:disabled) { background-color: #fb8c00; }
|
||||
.saturnnodes-btn-stop { background-color: #d32f2f; }
|
||||
.saturnnodes-btn-stop:hover:not(:disabled) { background-color: #f44336; }
|
||||
|
||||
.flowcontrol-decision-container {
|
||||
.saturnnodes-decision-container {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
@@ -37,46 +38,84 @@ styles.textContent = `
|
||||
`;
|
||||
document.head.appendChild(styles);
|
||||
|
||||
function findGraphNode(id) {
|
||||
if (!id && id !== 0) return null;
|
||||
const strId = String(id);
|
||||
const numId = Number(id);
|
||||
|
||||
// 1. Direct root graph lookup
|
||||
let node = app.graph?.getNodeById?.(id) || app.graph?.getNodeById?.(numId) || app.graph?.getNodeById?.(strId);
|
||||
if (node) return node;
|
||||
|
||||
// 2. Search root graph nodes and nested subgraphs
|
||||
if (app.graph && Array.isArray(app.graph._nodes)) {
|
||||
for (const n of app.graph._nodes) {
|
||||
if (String(n.id) === strId) return n;
|
||||
if (n.subgraph && typeof n.subgraph.getNodeById === "function") {
|
||||
const sub = n.subgraph.getNodeById(id) || n.subgraph.getNodeById(numId) || n.subgraph.getNodeById(strId);
|
||||
if (sub) return sub;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Search graph._subgraphs registry (v0.37+)
|
||||
if (app.graph && app.graph._subgraphs) {
|
||||
const subgraphs = Array.isArray(app.graph._subgraphs)
|
||||
? app.graph._subgraphs
|
||||
: Object.values(app.graph._subgraphs);
|
||||
for (const sg of subgraphs) {
|
||||
if (sg && typeof sg.getNodeById === "function") {
|
||||
const sub = sg.getNodeById(id) || sg.getNodeById(numId) || sg.getNodeById(strId);
|
||||
if (sub) return sub;
|
||||
}
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "Comfy.FlowControlDecision",
|
||||
name: "ComfyUI.SaturnNodes.Decision",
|
||||
async setup() {
|
||||
api.addEventListener("flowcontrol_decision_waiting", (event) => {
|
||||
const handleWaiting = (event) => {
|
||||
const data = event.detail;
|
||||
if (data && data.node_id) {
|
||||
const node = app.graph.getNodeById(data.node_id);
|
||||
const node = findGraphNode(data.node_id);
|
||||
if (node && node.enableDecisionButtons) {
|
||||
node.enableDecisionButtons();
|
||||
}
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
api.addEventListener("flowcontrol_decision_resolved", (event) => {
|
||||
const handleResolved = (event) => {
|
||||
const data = event.detail;
|
||||
if (data && data.node_id) {
|
||||
const node = app.graph.getNodeById(data.node_id);
|
||||
const node = findGraphNode(data.node_id);
|
||||
if (node && node.disableDecisionButtons) {
|
||||
node.disableDecisionButtons();
|
||||
}
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
api.addEventListener("saturnnodes_decision_waiting", handleWaiting);
|
||||
api.addEventListener("saturnnodes_decision_resolved", handleResolved);
|
||||
},
|
||||
async nodeCreated(node) {
|
||||
if (node.comfyClass === "FlowControlDecision") {
|
||||
if (node.comfyClass === "SaturnDecision") {
|
||||
const container = document.createElement("div");
|
||||
container.className = "flowcontrol-decision-container";
|
||||
container.className = "saturnnodes-decision-container";
|
||||
|
||||
const btnContinue = document.createElement("button");
|
||||
btnContinue.className = "flowcontrol-decision-btn flowcontrol-btn-continue";
|
||||
btnContinue.className = "saturnnodes-decision-btn saturnnodes-btn-continue";
|
||||
btnContinue.innerText = "Continue Workflow";
|
||||
btnContinue.disabled = true;
|
||||
|
||||
const btnCancel = document.createElement("button");
|
||||
btnCancel.className = "flowcontrol-decision-btn flowcontrol-btn-cancel";
|
||||
btnCancel.className = "saturnnodes-decision-btn saturnnodes-btn-cancel";
|
||||
btnCancel.innerText = "Cancel (Return True)";
|
||||
btnCancel.disabled = true;
|
||||
|
||||
const btnStop = document.createElement("button");
|
||||
btnStop.className = "flowcontrol-decision-btn flowcontrol-btn-stop";
|
||||
btnStop.className = "saturnnodes-decision-btn saturnnodes-btn-stop";
|
||||
btnStop.innerText = "Stop Workflow";
|
||||
btnStop.disabled = true;
|
||||
|
||||
@@ -87,15 +126,16 @@ app.registerExtension({
|
||||
const sendDecision = async (action) => {
|
||||
node.disableDecisionButtons();
|
||||
try {
|
||||
await api.fetchApi("/flowcontrol/decision", {
|
||||
await authenticatedFetch("/saturnnodes/decision", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
node_id: node.id.toString(),
|
||||
action: action
|
||||
}),
|
||||
});
|
||||
} catch (e) {
|
||||
console.error("[FlowControl] Failed to send decision:", e);
|
||||
console.error("[SaturnNodes] Failed to send decision:", e);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
+756
-143
File diff suppressed because it is too large
Load Diff
+146
-67
@@ -14,17 +14,10 @@ visualStyles.textContent = `
|
||||
border: 1px solid #222;
|
||||
border-radius: 6px;
|
||||
box-sizing: border-box;
|
||||
overflow-y: hidden;
|
||||
}
|
||||
.img-visual-container.vue-mode {
|
||||
height: 400px;
|
||||
min-height: 200px;
|
||||
resize: vertical;
|
||||
}
|
||||
.img-visual-container.litegraph-mode {
|
||||
height: 100%;
|
||||
min-height: 0;
|
||||
max-height: 100%;
|
||||
overflow-y: hidden;
|
||||
}
|
||||
.img-visual-container::-webkit-scrollbar {
|
||||
width: 6px;
|
||||
@@ -208,18 +201,18 @@ visualStyles.textContent = `
|
||||
document.head.appendChild(visualStyles);
|
||||
|
||||
app.registerExtension({
|
||||
name: "Comfy.ImageLoaderCustom",
|
||||
name: "ComfyUI.SaturnNodes.VisualImageLoader",
|
||||
async nodeCreated(node) {
|
||||
if (node.comfyClass === "ImageLoaderVisualPrettyV2" || node.comfyClass === "ImageLoaderCustom") {
|
||||
node.size = [380, 480];
|
||||
const isVisualImageLoader = [
|
||||
"VisualImageLoader",
|
||||
"ImageLoaderCustom",
|
||||
].includes(node.comfyClass) || [
|
||||
"VisualImageLoader",
|
||||
"ImageLoaderCustom",
|
||||
].includes(node.type);
|
||||
|
||||
node.computeSize = function() {
|
||||
const displayModeWidget = node.widgets ? node.widgets.find(w => w.name === "display_mode") : null;
|
||||
if (displayModeWidget && displayModeWidget.value === "Show All") {
|
||||
// Logic added below
|
||||
}
|
||||
return [node.size[0], Math.max(480, node.size[1])];
|
||||
};
|
||||
if (isVisualImageLoader) {
|
||||
node.size = [380, 480];
|
||||
|
||||
const folderWidget = node.widgets ? node.widgets.find(w => w.name === "folder_path") : null;
|
||||
|
||||
@@ -231,11 +224,13 @@ app.registerExtension({
|
||||
}
|
||||
}
|
||||
if (widgets.length === 1) {
|
||||
widgets[0].type = "hidden";
|
||||
widgets[0].hidden = true;
|
||||
widgets[0].computeSize = () => [0, 0];
|
||||
return widgets[0];
|
||||
}
|
||||
const w = node.addWidget("text", name, defaultVal, () => {}, { serialize: true });
|
||||
w.type = "hidden";
|
||||
w.hidden = true;
|
||||
w.computeSize = () => [0, 0];
|
||||
return w;
|
||||
@@ -253,11 +248,29 @@ app.registerExtension({
|
||||
|
||||
const controlBar = document.createElement("div");
|
||||
controlBar.className = "img-control-bar";
|
||||
controlBar.style.display = "flex";
|
||||
controlBar.style.flexDirection = "row";
|
||||
controlBar.style.gap = "6px";
|
||||
|
||||
const refreshBtn = document.createElement("button");
|
||||
refreshBtn.innerHTML = "↻";
|
||||
refreshBtn.title = "Refresh Images";
|
||||
refreshBtn.style.background = "#1c1c1c";
|
||||
refreshBtn.style.border = "1px solid #333";
|
||||
refreshBtn.style.borderRadius = "4px";
|
||||
refreshBtn.style.color = "#fff";
|
||||
refreshBtn.style.cursor = "pointer";
|
||||
refreshBtn.style.padding = "4px 8px";
|
||||
refreshBtn.style.fontSize = "14px";
|
||||
refreshBtn.onmouseover = () => refreshBtn.style.background = "#333";
|
||||
refreshBtn.onmouseout = () => refreshBtn.style.background = "#1c1c1c";
|
||||
controlBar.appendChild(refreshBtn);
|
||||
|
||||
const searchBar = document.createElement("input");
|
||||
searchBar.type = "text";
|
||||
searchBar.className = "img-search-bar";
|
||||
searchBar.placeholder = "Search images or positive prompts...";
|
||||
searchBar.style.flex = "1";
|
||||
controlBar.appendChild(searchBar);
|
||||
viewContainer.appendChild(controlBar);
|
||||
|
||||
@@ -265,50 +278,92 @@ app.registerExtension({
|
||||
gridContainer.className = "img-grid-container";
|
||||
viewContainer.appendChild(gridContainer);
|
||||
|
||||
node.computeSize = function() {
|
||||
const displayModeWidget = node.widgets ? node.widgets.find(w => w.name === "display_mode") : null;
|
||||
if (displayModeWidget && displayModeWidget.value === "Show All" && gridContainer) {
|
||||
return [node.size[0], Math.max(480, gridContainer.scrollHeight + 140)];
|
||||
}
|
||||
return [node.size[0], Math.max(480, node.size[1])];
|
||||
};
|
||||
let userSavedHeight = 480;
|
||||
|
||||
const updateNodeSize = () => {
|
||||
const displayModeWidget = node.widgets ? node.widgets.find(w => w.name === "display_mode") : null;
|
||||
if (displayModeWidget && displayModeWidget.value === "Show All") {
|
||||
const neededHeight = gridContainer.scrollHeight + 140;
|
||||
node.setSize([node.size[0], Math.max(480, neededHeight)]);
|
||||
const applyDisplayMode = (mode) => {
|
||||
const isShowAll = mode === "Show All";
|
||||
if (isShowAll) {
|
||||
viewContainer.style.setProperty("height", "auto", "important");
|
||||
viewContainer.style.setProperty("max-height", "none", "important");
|
||||
viewContainer.style.setProperty("overflow-y", "visible", "important");
|
||||
|
||||
gridContainer.style.setProperty("height", "auto", "important");
|
||||
gridContainer.style.setProperty("max-height", "none", "important");
|
||||
gridContainer.style.setProperty("overflow-y", "visible", "important");
|
||||
gridContainer.style.setProperty("flex", "none", "important");
|
||||
|
||||
setTimeout(() => {
|
||||
const contentH = gridContainer.scrollHeight || 0;
|
||||
const neededHeight = contentH + 160;
|
||||
node.setSize([node.size[0], Math.max(420, neededHeight)]);
|
||||
if (app.graph) app.graph.setDirtyCanvas(true, true);
|
||||
}, 50);
|
||||
} else {
|
||||
viewContainer.classList.remove("show-all-mode");
|
||||
viewContainer.classList.add("scrollable-mode");
|
||||
|
||||
viewContainer.style.setProperty("height", "100%", "important");
|
||||
viewContainer.style.setProperty("max-height", "100%", "important");
|
||||
viewContainer.style.setProperty("overflow-y", "hidden", "important");
|
||||
|
||||
gridContainer.style.setProperty("height", "100%", "important");
|
||||
gridContainer.style.setProperty("max-height", "100%", "important");
|
||||
gridContainer.style.setProperty("overflow-y", "auto", "important");
|
||||
gridContainer.style.setProperty("flex", "1 1 0px", "important");
|
||||
|
||||
const targetH = node.userCustomHeight || userSavedHeight || (node.size ? node.size[1] : 480);
|
||||
if (targetH && (!node.size || node.size[1] !== targetH)) {
|
||||
node.setSize([node.size ? node.size[0] : 380, Math.max(420, targetH)]);
|
||||
}
|
||||
if (app.graph) app.graph.setDirtyCanvas(true, true);
|
||||
}
|
||||
};
|
||||
|
||||
const initialZoom = localStorage.getItem("comfy_img_picker_zoom") || "80";
|
||||
node.computeSize = function() {
|
||||
const displayModeWidget = node.widgets ? node.widgets.find(w => w.name === "display_mode") : null;
|
||||
if (displayModeWidget && displayModeWidget.value === "Show All" && gridContainer) {
|
||||
return [Math.max(380, node.size[0]), Math.max(420, (gridContainer.scrollHeight || 0) + 160)];
|
||||
}
|
||||
return [Math.max(380, node.size[0]), Math.max(420, node.size ? node.size[1] : 420)];
|
||||
};
|
||||
|
||||
const updateNodeSize = () => {
|
||||
const displayModeWidget = node.widgets ? node.widgets.find(w => w.name === "display_mode") : null;
|
||||
const mode = displayModeWidget ? displayModeWidget.value : "Scrollable";
|
||||
applyDisplayMode(mode);
|
||||
};
|
||||
|
||||
node.properties = node.properties || {};
|
||||
const initialZoom = node.properties["tile_size"] != null
|
||||
? node.properties["tile_size"]
|
||||
: (localStorage.getItem("comfy_img_picker_zoom") || "80");
|
||||
node.properties["tile_size"] = parseInt(initialZoom);
|
||||
viewContainer.style.setProperty("--img-tile-size", `${initialZoom}px`);
|
||||
|
||||
// Embed DOM Widget inside the node container
|
||||
const domWidget = node.addDOMWidget("img_visual_picker", "HTML", viewContainer, {
|
||||
getValue() { return getHiddenWidget("_selected_image", "").value; },
|
||||
setValue(val) { getHiddenWidget("_selected_image", "").value = val; },
|
||||
serialize: false
|
||||
serialize: false,
|
||||
getMinHeight() { return 180; }
|
||||
});
|
||||
|
||||
// Detect Vue UI vs LiteGraph
|
||||
const isVueUI = !!document.querySelector(".lg-node");
|
||||
if (isVueUI) {
|
||||
viewContainer.classList.add("vue-mode");
|
||||
} else {
|
||||
viewContainer.classList.add("litegraph-mode");
|
||||
const onResize = node.onResize;
|
||||
node.onResize = function(size) {
|
||||
if (onResize) onResize.apply(this, arguments);
|
||||
if (viewContainer.style) {
|
||||
viewContainer.style.height = `${size[1] - 80}px`;
|
||||
}
|
||||
};
|
||||
}
|
||||
// In ComfyUI LiteGraph (V1), computeLayoutSize handles dynamic widget space distribution.
|
||||
// Removing computeSize avoids an infinite expansion loop in _arrangeWidgets.
|
||||
domWidget.computeLayoutSize = function() {
|
||||
return { minHeight: 180, maxHeight: undefined, minWidth: 300 };
|
||||
};
|
||||
delete domWidget.computeSize;
|
||||
|
||||
domWidget.computeSize = function() {
|
||||
return [node.size[0] - 30, node.size[1] - 140];
|
||||
const onResize = node.onResize;
|
||||
node.onResize = function(size) {
|
||||
if (onResize) onResize.apply(this, arguments);
|
||||
const displayModeWidget = node.widgets ? node.widgets.find(w => w.name === "display_mode") : null;
|
||||
const isShowAll = displayModeWidget && displayModeWidget.value === "Show All";
|
||||
if (!isShowAll && size) {
|
||||
node.userCustomHeight = size[1];
|
||||
userSavedHeight = size[1];
|
||||
}
|
||||
};
|
||||
|
||||
// Zoom Slider Widget
|
||||
@@ -317,8 +372,10 @@ app.registerExtension({
|
||||
"tile_size",
|
||||
parseInt(initialZoom),
|
||||
(val) => {
|
||||
localStorage.setItem("comfy_img_picker_zoom", val);
|
||||
node.properties = node.properties || {};
|
||||
node.properties["tile_size"] = val;
|
||||
viewContainer.style.setProperty("--img-tile-size", `${val}px`);
|
||||
if (app.graph) app.graph.setDirtyCanvas(true, true);
|
||||
},
|
||||
{ min: 50, max: 200, step: 1 }
|
||||
);
|
||||
@@ -348,8 +405,8 @@ app.registerExtension({
|
||||
gridContainer.innerHTML = "";
|
||||
lazyObserver.disconnect();
|
||||
|
||||
const folder = folderWidget ? (folderWidget.value || "") : "";
|
||||
const selectedVal = getHiddenWidget().value;
|
||||
let folder = folderWidget ? (folderWidget.value || "") : "";
|
||||
const selectedVal = getHiddenWidget("_selected_image", "").value;
|
||||
|
||||
// Group items by subfolder
|
||||
const foldersMap = {};
|
||||
@@ -450,7 +507,8 @@ app.registerExtension({
|
||||
}
|
||||
|
||||
const img = document.createElement("img");
|
||||
img.dataset.src = `/image_loader/get_thumbnail?folder=${encodeURIComponent(folder)}&image=${encodeURIComponent(item.relPath)}`;
|
||||
let fetchFolder = folder;
|
||||
img.dataset.src = `/saturnnodes/image_loader/get_thumbnail?folder=${encodeURIComponent(fetchFolder)}&image=${encodeURIComponent(item.relPath)}`;
|
||||
img.onload = () => img.classList.add("loaded");
|
||||
|
||||
img.onerror = () => {
|
||||
@@ -472,7 +530,7 @@ app.registerExtension({
|
||||
allTiles.forEach(t => t.classList.remove("selected"));
|
||||
|
||||
tile.classList.add("selected");
|
||||
const curWidget = getHiddenWidget();
|
||||
const curWidget = getHiddenWidget("_selected_image", "");
|
||||
curWidget.value = item.relPath;
|
||||
if (curWidget.callback) curWidget.callback(item.relPath);
|
||||
|
||||
@@ -496,12 +554,12 @@ app.registerExtension({
|
||||
};
|
||||
|
||||
const updateImagesList = async () => {
|
||||
const folder = folderWidget ? (folderWidget.value || "") : "";
|
||||
let folder = folderWidget ? (folderWidget.value || "") : "";
|
||||
const currentRequest = Symbol();
|
||||
activeRequest = currentRequest;
|
||||
|
||||
try {
|
||||
const response = await api.fetchApi(`/image_loader/get_images?folder=${encodeURIComponent(folder)}`);
|
||||
const response = await api.fetchApi(`/saturnnodes/image_loader/get_images?folder=${encodeURIComponent(folder)}`);
|
||||
if (activeRequest !== currentRequest) return;
|
||||
|
||||
imagesData = await response.json();
|
||||
@@ -517,6 +575,16 @@ app.registerExtension({
|
||||
debounceTimer = setTimeout(() => updateImagesList(), 350);
|
||||
};
|
||||
}
|
||||
|
||||
refreshBtn.onclick = () => {
|
||||
refreshBtn.style.transform = "rotate(180deg)";
|
||||
refreshBtn.style.transition = "transform 0.3s";
|
||||
updateImagesList();
|
||||
setTimeout(() => {
|
||||
refreshBtn.style.transform = "none";
|
||||
refreshBtn.style.transition = "none";
|
||||
}, 300);
|
||||
};
|
||||
|
||||
searchBar.addEventListener("input", () => {
|
||||
searchQuery = searchBar.value;
|
||||
@@ -528,7 +596,27 @@ app.registerExtension({
|
||||
node.onConfigure = function(config) {
|
||||
if (originalOnConfigure) originalOnConfigure.apply(this, arguments);
|
||||
|
||||
const curWidget = getHiddenWidget();
|
||||
if (config && config.size && config.size[1]) {
|
||||
node.userCustomHeight = config.size[1];
|
||||
userSavedHeight = config.size[1];
|
||||
}
|
||||
|
||||
const savedTileSize = (config && config.properties && config.properties["tile_size"] != null)
|
||||
? config.properties["tile_size"]
|
||||
: (node.properties && node.properties["tile_size"] != null ? node.properties["tile_size"] : null);
|
||||
|
||||
if (savedTileSize != null) {
|
||||
node.properties = node.properties || {};
|
||||
node.properties["tile_size"] = savedTileSize;
|
||||
if (zoomWidget) {
|
||||
zoomWidget.value = parseInt(savedTileSize);
|
||||
}
|
||||
if (viewContainer && viewContainer.style) {
|
||||
viewContainer.style.setProperty("--img-tile-size", `${savedTileSize}px`);
|
||||
}
|
||||
}
|
||||
|
||||
const curWidget = getHiddenWidget("_selected_image", "");
|
||||
const widgetIndex = node.widgets.indexOf(curWidget);
|
||||
const savedValue = (config.widgets_values && widgetIndex !== -1) ? config.widgets_values[widgetIndex] : curWidget.value;
|
||||
if (savedValue) {
|
||||
@@ -543,18 +631,9 @@ app.registerExtension({
|
||||
const origCallback = displayModeWidget.callback;
|
||||
displayModeWidget.callback = function(val) {
|
||||
if (origCallback) origCallback.apply(this, arguments);
|
||||
if (val === "Show All") {
|
||||
gridContainer.style.overflowY = "hidden";
|
||||
} else {
|
||||
gridContainer.style.overflowY = "auto";
|
||||
}
|
||||
setTimeout(() => updateNodeSize(), 50);
|
||||
applyDisplayMode(val || displayModeWidget.value);
|
||||
};
|
||||
if (displayModeWidget.value === "Show All") {
|
||||
gridContainer.style.overflowY = "hidden";
|
||||
} else {
|
||||
gridContainer.style.overflowY = "auto";
|
||||
}
|
||||
applyDisplayMode(displayModeWidget.value || "Scrollable");
|
||||
}
|
||||
updateImagesList();
|
||||
}, 100);
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
/**
|
||||
* Helpers for DOM widgets whose elements can be unmounted and rebuilt by the
|
||||
* ComfyUI Vue node renderer or LiteGraph. A cached element reference then
|
||||
* points at a detached subtree, so writes silently update an invisible node.
|
||||
*/
|
||||
|
||||
export function isLiveElement(el) {
|
||||
return !!el && (typeof el.isConnected !== "boolean" || el.isConnected);
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the cached element while it is still live. When it is detached, try
|
||||
* to find the element the renderer remounted for the same node/widget.
|
||||
*
|
||||
* Before the first mount there is no replacement yet, so the cached reference
|
||||
* is returned to preserve initial-value writes.
|
||||
*/
|
||||
export function resolveLiveElement(cachedEl, findFresh) {
|
||||
if (isLiveElement(cachedEl)) return cachedEl;
|
||||
if (typeof findFresh === "function") {
|
||||
try {
|
||||
const fresh = findFresh();
|
||||
if (fresh) return fresh;
|
||||
} catch (_) {}
|
||||
}
|
||||
return cachedEl || null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Re-resolve a node DOM widget by its `data-widget-name` marker.
|
||||
* Used by the prompt counter and LoRA count badges.
|
||||
*/
|
||||
export function resolveLiveWidgetElement(cachedEl, nodeId, widgetName) {
|
||||
return resolveLiveElement(cachedEl, () => {
|
||||
if (nodeId === undefined || nodeId === null) return null;
|
||||
const id = String(nodeId);
|
||||
return document.querySelector(
|
||||
`[data-node-id="${id}"] [data-widget-name="${widgetName}"], ` +
|
||||
`[data-widgets-grid-node-id="${id}"] [data-widget-name="${widgetName}"]`
|
||||
);
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { authenticatedFetch } from "./auth_helper.js";
|
||||
|
||||
// SaturnNodes Local File Runner In-Memory Consent Gate Extension
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.SaturnNodes.LocalRunner",
|
||||
async nodeCreated(node) {
|
||||
if (node.comfyClass !== "RunLocalFileNode") return;
|
||||
|
||||
node._isAuthorized = false;
|
||||
node._authExpiresAt = 0;
|
||||
|
||||
const filePathWidget = node.widgets?.find(w => w.name === "file_path");
|
||||
|
||||
const updateVisualStatus = () => {
|
||||
const now = Date.now();
|
||||
const isArmed = Boolean(node._isAuthorized && node._authExpiresAt > now);
|
||||
|
||||
if (isArmed) {
|
||||
const remainingSec = Math.max(0, Math.round((node._authExpiresAt - now) / 1000));
|
||||
const remainingMin = Math.ceil(remainingSec / 60);
|
||||
node.title = `🪐 ⚡ Run Local File [ARMED (${remainingMin}m)]`;
|
||||
node.color = "#047857";
|
||||
node.bgcolor = "#064e3b";
|
||||
if (authButton) {
|
||||
authButton.name = `✅ Authorized (${remainingMin}m left)`;
|
||||
}
|
||||
} else {
|
||||
node._isAuthorized = false;
|
||||
node.title = "🪐 ⚡ Run Local File [DISARMED]";
|
||||
node.color = "#7f1d1d";
|
||||
node.bgcolor = "#18181b";
|
||||
if (authButton) {
|
||||
authButton.name = "⚡ Authorize Run (5 min)";
|
||||
}
|
||||
}
|
||||
node.setDirtyCanvas(true, true);
|
||||
};
|
||||
|
||||
// Add explicit interactive action button: "⚡ Authorize Run (5 min)"
|
||||
const authButton = node.addWidget(
|
||||
"button",
|
||||
"⚡ Authorize Run (5 min)",
|
||||
null,
|
||||
async () => {
|
||||
const scriptName = filePathWidget?.value?.trim() || "configured script";
|
||||
const confirmed = confirm(
|
||||
"🪐 SaturnNodes Security Verification\n\n" +
|
||||
`Authorize a SINGLE run of local script:\n'${scriptName}'\n\n` +
|
||||
"Notice: Scripts can execute shell commands on your system.\n" +
|
||||
"• The script must reside inside 'ComfyUI/scripts/'.\n" +
|
||||
"• Authorization is temporary (expires in 5 minutes) and single-use.\n" +
|
||||
"• External workflows can NEVER run scripts without your live authorization.\n\n" +
|
||||
"Authorize this run now?"
|
||||
);
|
||||
|
||||
if (!confirmed) return;
|
||||
|
||||
try {
|
||||
const resp = await authenticatedFetch("/saturnnodes/local_runner/authorize", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
node_id: node.id,
|
||||
file_path: filePathWidget?.value || ""
|
||||
})
|
||||
});
|
||||
|
||||
if (!resp.ok) {
|
||||
const err = await resp.json().catch(() => ({}));
|
||||
const errMsg = err.error || "Failed to authorize node run.";
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🪐 Authorization Failed",
|
||||
detail: errMsg,
|
||||
life: 6000
|
||||
});
|
||||
} else {
|
||||
alert("❌ " + errMsg);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
node._isAuthorized = true;
|
||||
node._authExpiresAt = Date.now() + (300 * 1000);
|
||||
updateVisualStatus();
|
||||
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "success",
|
||||
summary: "⚡ Script Run Authorized",
|
||||
detail: `Node ${node.id} authorized for 5 minutes. Single-use: will be consumed upon execution.`,
|
||||
life: 5000
|
||||
});
|
||||
}
|
||||
} catch (e) {
|
||||
console.error("[SaturnNodes Security] Error authorizing script run:", e);
|
||||
alert("❌ Error communicating with authorization service: " + e.message);
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
// Intercept workflow load / configuration: ALWAYS disarm on load!
|
||||
const origOnConfigure = node.onConfigure;
|
||||
node.onConfigure = function(info) {
|
||||
if (origOnConfigure) origOnConfigure.apply(this, arguments);
|
||||
// GUARANTEE: Never let imported workflow or image auto-run with consent pre-checked
|
||||
node._isAuthorized = false;
|
||||
node._authExpiresAt = 0;
|
||||
updateVisualStatus();
|
||||
};
|
||||
|
||||
// If file path changes, invalidate existing authorization
|
||||
if (filePathWidget) {
|
||||
const origCallback = filePathWidget.callback;
|
||||
filePathWidget.callback = function(val) {
|
||||
node._isAuthorized = false;
|
||||
node._authExpiresAt = 0;
|
||||
updateVisualStatus();
|
||||
if (origCallback) origCallback.apply(this, arguments);
|
||||
};
|
||||
}
|
||||
|
||||
// Draw security alert ribbon on canvas if disarmed
|
||||
const origOnDrawForeground = node.onDrawForeground;
|
||||
node.onDrawForeground = function(ctx) {
|
||||
if (origOnDrawForeground) origOnDrawForeground.apply(this, arguments);
|
||||
if (node.flags?.collapsed) return;
|
||||
|
||||
const now = Date.now();
|
||||
const isArmed = Boolean(node._isAuthorized && node._authExpiresAt > now);
|
||||
|
||||
ctx.save();
|
||||
ctx.textAlign = "center";
|
||||
ctx.font = "italic 9px sans-serif";
|
||||
ctx.fillStyle = "#9ca3af";
|
||||
ctx.fillText("Script Location: ComfyUI/scripts/", node.size[0] / 2, node.size[1] - 19);
|
||||
|
||||
ctx.font = "bold 10px sans-serif";
|
||||
if (!isArmed) {
|
||||
ctx.fillStyle = "#ef4444";
|
||||
ctx.fillText("🛡️ DISARMED: Click 'Authorize Run' to execute", node.size[0] / 2, node.size[1] - 7);
|
||||
} else {
|
||||
const rem = Math.max(0, Math.ceil((node._authExpiresAt - now) / 60000));
|
||||
ctx.fillStyle = "#10b981";
|
||||
ctx.fillText(`⚡ ARMED (${rem}m left) - Consumed on run`, node.size[0] / 2, node.size[1] - 7);
|
||||
}
|
||||
ctx.restore();
|
||||
};
|
||||
|
||||
updateVisualStatus();
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,718 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
import { authenticatedFetch } from "./auth_helper.js";
|
||||
|
||||
const EXTENSION_NAME = "SaturnNodes.PowerControlSidebar";
|
||||
const STYLE_ID = "saturnnodes-power-sidebar-styles";
|
||||
|
||||
let currentPowerState = {
|
||||
pending_action: null,
|
||||
armed_at: null,
|
||||
is_paused: false
|
||||
};
|
||||
|
||||
let activePopup = null;
|
||||
let activePopupCleanup = null;
|
||||
|
||||
function injectStyles() {
|
||||
if (document.getElementById(STYLE_ID)) return;
|
||||
const style = document.createElement("style");
|
||||
style.id = STYLE_ID;
|
||||
style.textContent = `
|
||||
/* Power Sidebar Button Layout (Matching standard full-width sidebar button) */
|
||||
#saturnnodes-power-bottom-btn {
|
||||
width: 100% !important;
|
||||
height: auto !important;
|
||||
min-height: 48px !important;
|
||||
padding: 6px 4px !important;
|
||||
box-sizing: border-box !important;
|
||||
border-radius: 8px !important;
|
||||
align-self: stretch !important;
|
||||
}
|
||||
|
||||
#saturnnodes-power-bottom-btn .side-bar-button-content {
|
||||
width: 100% !important;
|
||||
display: flex !important;
|
||||
flex-direction: column !important;
|
||||
align-items: center !important;
|
||||
gap: 4px !important;
|
||||
}
|
||||
|
||||
/* Power Sidebar Button Armed Red Warning State */
|
||||
#saturnnodes-power-bottom-btn.saturnnodes-power-btn-armed,
|
||||
.comfy-power-btn.saturnnodes-power-btn-armed,
|
||||
.side-bar-button.saturnnodes-power-btn-armed,
|
||||
.side-bar-button:has(.saturnnodes-power-icon-armed),
|
||||
button[aria-label*="Power"].saturnnodes-power-btn-armed {
|
||||
background: linear-gradient(135deg, #dc2626 0%, #991b1b 100%) !important;
|
||||
color: #ffffff !important;
|
||||
border: 1px solid #f87171 !important;
|
||||
box-shadow: 0 0 14px rgba(239, 68, 68, 0.8) !important;
|
||||
animation: saturnnodes-power-pulse 1.8s infinite ease-in-out !important;
|
||||
}
|
||||
|
||||
.saturnnodes-power-btn-armed svg,
|
||||
.saturnnodes-power-btn-armed .side-bar-button-icon,
|
||||
.saturnnodes-power-btn-armed span {
|
||||
color: #ffffff !important;
|
||||
stroke: #ffffff !important;
|
||||
}
|
||||
|
||||
@keyframes saturnnodes-power-pulse {
|
||||
0% { box-shadow: 0 0 6px rgba(239, 68, 68, 0.5); transform: scale(1); }
|
||||
50% { box-shadow: 0 0 16px rgba(239, 68, 68, 0.95); transform: scale(1.03); }
|
||||
100% { box-shadow: 0 0 6px rgba(239, 68, 68, 0.5); transform: scale(1); }
|
||||
}
|
||||
|
||||
/* Power Side Popup Menu (like Help Center popup) */
|
||||
.saturnnodes-power-popup {
|
||||
position: fixed;
|
||||
z-index: 10005;
|
||||
min-width: 250px;
|
||||
max-width: 320px;
|
||||
background: var(--comfy-menu-bg, #18181b);
|
||||
border: 1px solid var(--border-color, #27272a);
|
||||
border-radius: 12px;
|
||||
box-shadow: 0 12px 30px rgba(0, 0, 0, 0.6);
|
||||
padding: 8px;
|
||||
color: var(--fg-color, #f4f4f5);
|
||||
font-family: inherit;
|
||||
animation: saturnnodes-popup-fade 0.15s ease-out;
|
||||
}
|
||||
|
||||
@keyframes saturnnodes-popup-fade {
|
||||
from { opacity: 0; transform: translateX(-6px); }
|
||||
to { opacity: 1; transform: translateX(0); }
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .help-menu-item {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 10px;
|
||||
width: 100%;
|
||||
padding: 8px 10px;
|
||||
border-radius: 8px;
|
||||
border: none;
|
||||
background: transparent;
|
||||
color: var(--fg-color, #f4f4f5);
|
||||
font-size: 13px;
|
||||
font-weight: 500;
|
||||
cursor: pointer;
|
||||
text-align: left;
|
||||
transition: background 0.15s ease, color 0.15s ease;
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .help-menu-item:hover {
|
||||
background: var(--comfy-menu-secondary-bg, rgba(255, 255, 255, 0.08));
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .help-menu-item.danger-item:hover {
|
||||
background: rgba(239, 68, 68, 0.15);
|
||||
color: #f87171;
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .help-menu-item.armed-cancel-item {
|
||||
background: rgba(239, 68, 68, 0.85);
|
||||
color: #ffffff;
|
||||
font-weight: 600;
|
||||
margin-bottom: 6px;
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .help-menu-item.armed-cancel-item:hover {
|
||||
background: #dc2626;
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .help-menu-icon-container {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
width: 24px;
|
||||
height: 24px;
|
||||
flex-shrink: 0;
|
||||
color: var(--descrip-text, #a1a1aa);
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .help-menu-item:hover .help-menu-icon-container {
|
||||
color: var(--fg-color, #f4f4f5);
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .menu-label {
|
||||
flex: 1;
|
||||
white-space: nowrap;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
}
|
||||
|
||||
.saturnnodes-power-popup .menu-divider {
|
||||
height: 1px;
|
||||
background: var(--border-color, #27272a);
|
||||
margin: 4px 0;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
/* Power Confirm Modal */
|
||||
.saturnnodes-power-modal {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 12px;
|
||||
padding: 18px;
|
||||
min-width: 330px;
|
||||
max-width: 440px;
|
||||
background: #18181b;
|
||||
color: #f4f4f5;
|
||||
border-radius: 12px;
|
||||
font-family: inherit;
|
||||
}
|
||||
`;
|
||||
document.head.appendChild(style);
|
||||
}
|
||||
|
||||
function getPowerIconSvg(size = 18) {
|
||||
return `<svg class="size-4.5" width="${size}" height="${size}" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M18.36 6.64a9 9 0 1 1-12.73 0"></path><line x1="12" y1="2" x2="12" y2="12"></line></svg>`;
|
||||
}
|
||||
|
||||
function updateSidebarButtonVisuals() {
|
||||
const isArmed = Boolean(currentPowerState.pending_action);
|
||||
const btns = document.querySelectorAll('#saturnnodes-power-bottom-btn, .comfy-power-btn, [aria-label*="Power"], .side-bar-button:has(.pi-power)');
|
||||
btns.forEach(btn => {
|
||||
if (isArmed) {
|
||||
btn.classList.add("saturnnodes-power-btn-armed");
|
||||
const icon = btn.querySelector(".side-bar-button-icon, i, svg");
|
||||
if (icon) icon.classList.add("saturnnodes-power-icon-armed");
|
||||
btn.setAttribute("title", `⚠️ ComfyUI Power: Armed to ${currentPowerState.pending_action.toUpperCase()} after queue finish! Click to cancel.`);
|
||||
} else {
|
||||
btn.classList.remove("saturnnodes-power-btn-armed");
|
||||
const icon = btn.querySelector(".side-bar-button-icon, i, svg");
|
||||
if (icon) icon.classList.remove("saturnnodes-power-icon-armed");
|
||||
btn.setAttribute("title", "Power (Restart / Shutdown)");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async function fetchPowerStatus() {
|
||||
try {
|
||||
const resp = await api.fetchApi("/saturnnodes/power/status");
|
||||
if (resp.ok) {
|
||||
currentPowerState = await resp.json();
|
||||
updateSidebarButtonVisuals();
|
||||
}
|
||||
} catch (_) {}
|
||||
}
|
||||
|
||||
function openConfirmModal(title, message, onConfirm, isDanger = false) {
|
||||
const overlay = document.createElement("div");
|
||||
overlay.style.position = "fixed";
|
||||
overlay.style.inset = "0";
|
||||
overlay.style.backgroundColor = "rgba(0,0,0,0.75)";
|
||||
overlay.style.display = "flex";
|
||||
overlay.style.alignItems = "center";
|
||||
overlay.style.justifyContent = "center";
|
||||
overlay.style.zIndex = "10010";
|
||||
overlay.style.backdropFilter = "blur(3px)";
|
||||
|
||||
const modal = document.createElement("div");
|
||||
modal.className = "saturnnodes-power-modal";
|
||||
modal.style.border = isDanger ? "1px solid rgba(239, 68, 68, 0.5)" : "1px solid rgba(59, 130, 246, 0.5)";
|
||||
modal.style.boxShadow = isDanger ? "0 0 25px rgba(239, 68, 68, 0.35)" : "0 0 25px rgba(59, 130, 246, 0.35)";
|
||||
|
||||
const header = document.createElement("div");
|
||||
header.style.fontSize = "16px";
|
||||
header.style.fontWeight = "bold";
|
||||
header.style.display = "flex";
|
||||
header.style.alignItems = "center";
|
||||
header.style.gap = "8px";
|
||||
header.innerHTML = isDanger ? `⚠️ ${title}` : `ℹ️ ${title}`;
|
||||
|
||||
const desc = document.createElement("div");
|
||||
desc.style.fontSize = "13px";
|
||||
desc.style.lineHeight = "1.5";
|
||||
desc.style.color = "#d4d4d8";
|
||||
desc.textContent = message;
|
||||
|
||||
const btnRow = document.createElement("div");
|
||||
btnRow.style.display = "flex";
|
||||
btnRow.style.justifyContent = "flex-end";
|
||||
btnRow.style.gap = "8px";
|
||||
btnRow.style.marginTop = "8px";
|
||||
|
||||
const cancelBtn = document.createElement("button");
|
||||
cancelBtn.textContent = onConfirm ? "Cancel" : "Close";
|
||||
cancelBtn.style.padding = "8px 16px";
|
||||
cancelBtn.style.borderRadius = "6px";
|
||||
cancelBtn.style.border = "1px solid #52525b";
|
||||
cancelBtn.style.background = "#27272a";
|
||||
cancelBtn.style.color = "#f4f4f5";
|
||||
cancelBtn.style.cursor = "pointer";
|
||||
cancelBtn.onclick = () => overlay.remove();
|
||||
|
||||
if (onConfirm) {
|
||||
const confirmBtn = document.createElement("button");
|
||||
confirmBtn.textContent = "Confirm";
|
||||
confirmBtn.style.padding = "8px 16px";
|
||||
confirmBtn.style.borderRadius = "6px";
|
||||
confirmBtn.style.border = "none";
|
||||
confirmBtn.style.fontWeight = "600";
|
||||
confirmBtn.style.cursor = "pointer";
|
||||
if (isDanger) {
|
||||
confirmBtn.style.background = "#dc2626";
|
||||
confirmBtn.style.color = "#ffffff";
|
||||
} else {
|
||||
confirmBtn.style.background = "#2563eb";
|
||||
confirmBtn.style.color = "#ffffff";
|
||||
}
|
||||
|
||||
confirmBtn.onclick = async () => {
|
||||
overlay.remove();
|
||||
await onConfirm();
|
||||
};
|
||||
|
||||
btnRow.append(cancelBtn, confirmBtn);
|
||||
} else {
|
||||
btnRow.append(cancelBtn);
|
||||
}
|
||||
modal.append(header, desc, btnRow);
|
||||
overlay.appendChild(modal);
|
||||
document.body.appendChild(overlay);
|
||||
}
|
||||
|
||||
function checkProcessManagementEnabled() {
|
||||
if (currentPowerState.enabled === false) {
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "warn",
|
||||
summary: "🪐 Process Management Disabled",
|
||||
detail: "Enable 'Allow Process Management' in SaturnNodes settings to restart or shut down ComfyUI.",
|
||||
life: 7000
|
||||
});
|
||||
}
|
||||
openConfirmModal(
|
||||
"Process Management Disabled",
|
||||
"Server restart and shutdown controls are currently disabled for security.\n\nTo enable them, go to ComfyUI Settings ⚙️ -> SaturnNodes -> Pause Controls -> Turn ON 'Allow Process Management (Restart / Shutdown)'.",
|
||||
null,
|
||||
false
|
||||
);
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
function closePowerPopup() {
|
||||
// Document-level listeners must die with the popup they belong to.
|
||||
// Otherwise an old outside-click handler can close a popup that was
|
||||
// reopened after this one was dismissed by a menu item.
|
||||
if (activePopupCleanup) {
|
||||
try {
|
||||
activePopupCleanup();
|
||||
} catch (_) {}
|
||||
activePopupCleanup = null;
|
||||
}
|
||||
if (activePopup) {
|
||||
activePopup.remove();
|
||||
activePopup = null;
|
||||
}
|
||||
}
|
||||
|
||||
async function powerFetch(endpointSuffix, options = {}) {
|
||||
return authenticatedFetch(`/saturnnodes/power/${endpointSuffix}`, options);
|
||||
}
|
||||
|
||||
function showPowerSidePopup(anchorElement) {
|
||||
if (activePopup && activePopup.isConnected) {
|
||||
closePowerPopup();
|
||||
return;
|
||||
}
|
||||
if (activePopup) {
|
||||
// Detached popup left over from a DOM swap: clean it up before opening.
|
||||
closePowerPopup();
|
||||
}
|
||||
|
||||
const popup = document.createElement("div");
|
||||
popup.className = "help-center-popup sidebar-left saturnnodes-power-popup";
|
||||
popup.setAttribute("data-testid", "power-popup");
|
||||
|
||||
const rect = anchorElement.getBoundingClientRect();
|
||||
popup.style.left = `${Math.round(rect.right + 8)}px`;
|
||||
popup.style.bottom = `${Math.max(12, Math.round(window.innerHeight - rect.bottom))}px`;
|
||||
|
||||
const menu = document.createElement("div");
|
||||
menu.className = "help-center-menu flex flex-col items-start gap-1";
|
||||
menu.setAttribute("role", "menu");
|
||||
menu.setAttribute("aria-label", "Power Menu");
|
||||
|
||||
const wrap = document.createElement("div");
|
||||
wrap.className = "w-full";
|
||||
|
||||
const nav = document.createElement("nav");
|
||||
nav.className = "flex w-full flex-col gap-1";
|
||||
nav.setAttribute("role", "menubar");
|
||||
|
||||
// Cancel Armed Action Button (if armed)
|
||||
if (currentPowerState.pending_action) {
|
||||
const cancelBtn = document.createElement("button");
|
||||
cancelBtn.type = "button";
|
||||
cancelBtn.className = "help-menu-item armed-cancel-item";
|
||||
cancelBtn.setAttribute("role", "menuitem");
|
||||
cancelBtn.innerHTML = `
|
||||
<div class="help-menu-icon-container"><svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><circle cx="12" cy="12" r="10"></circle><line x1="15" y1="9" x2="9" y2="15"></line><line x1="9" y1="9" x2="15" y2="15"></line></svg></div>
|
||||
<span class="menu-label">Cancel ${currentPowerState.pending_action.toUpperCase()}</span>
|
||||
`;
|
||||
cancelBtn.onclick = async (e) => {
|
||||
e.stopPropagation();
|
||||
closePowerPopup();
|
||||
await powerFetch("arm", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ action: null })
|
||||
});
|
||||
await fetchPowerStatus();
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "info",
|
||||
summary: "🪐 SaturnNodes Power",
|
||||
detail: "Scheduled power action has been cancelled.",
|
||||
life: 3000
|
||||
});
|
||||
}
|
||||
};
|
||||
nav.appendChild(cancelBtn);
|
||||
}
|
||||
|
||||
// 1. Restart Server (Immediate)
|
||||
const restartItem = document.createElement("button");
|
||||
restartItem.type = "button";
|
||||
restartItem.className = "help-menu-item";
|
||||
restartItem.setAttribute("role", "menuitem");
|
||||
restartItem.innerHTML = `
|
||||
<div class="help-menu-icon-container"><svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><polyline points="23 4 23 10 17 10"></polyline><polyline points="1 20 1 14 7 14"></polyline><path d="M3.51 9a9 9 0 0 1 14.85-3.36L23 10M1 14l4.64 4.36A9 9 0 0 0 20.49 15"></path></svg></div>
|
||||
<span class="menu-label">Restart Server (Immediate)</span>
|
||||
`;
|
||||
restartItem.onclick = (e) => {
|
||||
e.stopPropagation();
|
||||
closePowerPopup();
|
||||
if (!checkProcessManagementEnabled()) return;
|
||||
openConfirmModal(
|
||||
"Restart ComfyUI Immediately?",
|
||||
"Are you sure you want to restart ComfyUI immediately? Any active prompt generation will be interrupted.",
|
||||
async () => {
|
||||
const tokenResp = await powerFetch("request_token", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ action: "restart" })
|
||||
});
|
||||
if (!tokenResp.ok) {
|
||||
const err = await tokenResp.json().catch(() => ({}));
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🪐 Restart Denied",
|
||||
detail: err.error || "Failed to obtain power action ticket.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
const tokenData = await tokenResp.json();
|
||||
const resp = await powerFetch("confirm_action", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ ticket: tokenData.ticket, action: "restart" })
|
||||
});
|
||||
if (!resp.ok) {
|
||||
const err = await resp.json().catch(() => ({}));
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🪐 Restart Failed",
|
||||
detail: err.error || "Failed to trigger restart.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "warn",
|
||||
summary: "🔄 Restarting ComfyUI",
|
||||
detail: "The server process is restarting. The UI will reload shortly.",
|
||||
life: 8000
|
||||
});
|
||||
}
|
||||
setTimeout(() => window.location.reload(), 2500);
|
||||
},
|
||||
true
|
||||
);
|
||||
};
|
||||
nav.appendChild(restartItem);
|
||||
|
||||
// 2. Shutdown Server (Immediate)
|
||||
const shutdownItem = document.createElement("button");
|
||||
shutdownItem.type = "button";
|
||||
shutdownItem.className = "help-menu-item danger-item";
|
||||
shutdownItem.setAttribute("role", "menuitem");
|
||||
shutdownItem.innerHTML = `
|
||||
<div class="help-menu-icon-container"><svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M18.36 6.64a9 9 0 1 1-12.73 0"></path><line x1="12" y1="2" x2="12" y2="12"></line></svg></div>
|
||||
<span class="menu-label">Shutdown Server (Immediate)</span>
|
||||
`;
|
||||
shutdownItem.onclick = (e) => {
|
||||
e.stopPropagation();
|
||||
closePowerPopup();
|
||||
if (!checkProcessManagementEnabled()) return;
|
||||
openConfirmModal(
|
||||
"Shutdown ComfyUI Server?",
|
||||
"Are you sure you want to SHUT DOWN the ComfyUI server immediately? The process will exit.",
|
||||
async () => {
|
||||
const tokenResp = await powerFetch("request_token", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ action: "shutdown" })
|
||||
});
|
||||
if (!tokenResp.ok) {
|
||||
const err = await tokenResp.json().catch(() => ({}));
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🪐 Shutdown Denied",
|
||||
detail: err.error || "Failed to obtain power action ticket.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
const tokenData = await tokenResp.json();
|
||||
const resp = await powerFetch("confirm_action", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ ticket: tokenData.ticket, action: "shutdown" })
|
||||
});
|
||||
if (!resp.ok) {
|
||||
const err = await resp.json().catch(() => ({}));
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🪐 Shutdown Failed",
|
||||
detail: err.error || "Failed to trigger shutdown.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🛑 ComfyUI Shutting Down",
|
||||
detail: "The server process is terminating.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
},
|
||||
true
|
||||
);
|
||||
};
|
||||
nav.appendChild(shutdownItem);
|
||||
|
||||
// Divider
|
||||
const divider = document.createElement("div");
|
||||
divider.className = "menu-divider";
|
||||
nav.appendChild(divider);
|
||||
|
||||
// 3. Restart After Queue Finish
|
||||
const restartQueueItem = document.createElement("button");
|
||||
restartQueueItem.type = "button";
|
||||
restartQueueItem.className = "help-menu-item";
|
||||
restartQueueItem.setAttribute("role", "menuitem");
|
||||
restartQueueItem.innerHTML = `
|
||||
<div class="help-menu-icon-container"><svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><circle cx="12" cy="12" r="10"></circle><polyline points="12 6 12 12 16 14"></polyline></svg></div>
|
||||
<span class="menu-label">Restart After Queue Finish</span>
|
||||
`;
|
||||
restartQueueItem.onclick = (e) => {
|
||||
e.stopPropagation();
|
||||
closePowerPopup();
|
||||
if (!checkProcessManagementEnabled()) return;
|
||||
openConfirmModal(
|
||||
"Schedule Restart After Queue Finish?",
|
||||
"ComfyUI will wait until all remaining prompts in the queue finish and the system returns to idle before restarting. (Note: If queue is paused, it will not restart).",
|
||||
async () => {
|
||||
const resp = await powerFetch("arm", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ action: "restart" })
|
||||
});
|
||||
if (!resp.ok) {
|
||||
const err = await resp.json().catch(() => ({}));
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🪐 Scheduling Failed",
|
||||
detail: err.error || "Failed to schedule restart.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
await fetchPowerStatus();
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "warn",
|
||||
summary: "⏳ Restart Scheduled",
|
||||
detail: "ComfyUI will restart automatically when the queue finishes.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
}
|
||||
);
|
||||
};
|
||||
nav.appendChild(restartQueueItem);
|
||||
|
||||
// 4. Shutdown After Queue Finish
|
||||
const shutdownQueueItem = document.createElement("button");
|
||||
shutdownQueueItem.type = "button";
|
||||
shutdownQueueItem.className = "help-menu-item danger-item";
|
||||
shutdownQueueItem.setAttribute("role", "menuitem");
|
||||
shutdownQueueItem.innerHTML = `
|
||||
<div class="help-menu-icon-container"><svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M21 12.79A9 9 0 1 1 11.21 3 7 7 0 0 0 21 12.79z"></path></svg></div>
|
||||
<span class="menu-label">Shutdown After Queue Finish</span>
|
||||
`;
|
||||
shutdownQueueItem.onclick = (e) => {
|
||||
e.stopPropagation();
|
||||
closePowerPopup();
|
||||
if (!checkProcessManagementEnabled()) return;
|
||||
openConfirmModal(
|
||||
"Schedule Shutdown After Queue Finish?",
|
||||
"ComfyUI will wait until all queued prompts complete and the system returns to idle before shutting down completely.",
|
||||
async () => {
|
||||
const resp = await powerFetch("arm", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ action: "shutdown" })
|
||||
});
|
||||
if (!resp.ok) {
|
||||
const err = await resp.json().catch(() => ({}));
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "🪐 Scheduling Failed",
|
||||
detail: err.error || "Failed to schedule shutdown.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
await fetchPowerStatus();
|
||||
if (app.extensionManager?.toast?.add) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: "warn",
|
||||
summary: "🌙 Shutdown Scheduled",
|
||||
detail: "ComfyUI will shut down automatically when the queue finishes.",
|
||||
life: 6000
|
||||
});
|
||||
}
|
||||
}
|
||||
);
|
||||
};
|
||||
nav.appendChild(shutdownQueueItem);
|
||||
|
||||
wrap.appendChild(nav);
|
||||
menu.appendChild(wrap);
|
||||
popup.appendChild(menu);
|
||||
document.body.appendChild(popup);
|
||||
activePopup = popup;
|
||||
|
||||
// Close on click outside or escape
|
||||
const outsideListener = (e) => {
|
||||
if (popup && !popup.contains(e.target) && !anchorElement.contains(e.target)) {
|
||||
closePowerPopup();
|
||||
}
|
||||
};
|
||||
const keyListener = (e) => {
|
||||
if (e.key === "Escape") {
|
||||
closePowerPopup();
|
||||
}
|
||||
};
|
||||
const removePopupListeners = () => {
|
||||
document.removeEventListener("pointerdown", outsideListener);
|
||||
document.removeEventListener("keydown", keyListener);
|
||||
};
|
||||
activePopupCleanup = removePopupListeners;
|
||||
|
||||
setTimeout(() => {
|
||||
// The popup may have been closed (or replaced) before the timer fires.
|
||||
if (activePopup !== popup) return;
|
||||
document.addEventListener("pointerdown", outsideListener);
|
||||
document.addEventListener("keydown", keyListener);
|
||||
}, 10);
|
||||
}
|
||||
|
||||
function attachBottomSidebarButton() {
|
||||
// Fast path for the DOM MutationObserver: bail out before any querySelector
|
||||
// work when our button is already mounted (this callback fires for every
|
||||
// mutation in the app while dragging/resizing UI elements).
|
||||
const existingBtn = document.getElementById("saturnnodes-power-bottom-btn");
|
||||
if (existingBtn && existingBtn.isConnected) return;
|
||||
|
||||
const bottomGroup = document.querySelector(".sidebar-item-group.mt-auto, .side-bar .mt-auto, [data-testid='bottom-sidebar']");
|
||||
if (!bottomGroup) return;
|
||||
|
||||
const btn = document.createElement("button");
|
||||
btn.id = "saturnnodes-power-bottom-btn";
|
||||
btn.type = "button";
|
||||
btn.className = "relative inline-flex items-center justify-center gap-2 touch-manipulation whitespace-nowrap appearance-none font-medium font-inter transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none [&_svg:not([width]):not([height])]:size-4 [&_svg]:shrink-0 bg-transparent text-muted-foreground hover:bg-secondary-background-hover h-8 rounded-lg p-2 text-xs side-bar-button cursor-pointer border-none comfy-power-btn";
|
||||
btn.setAttribute("aria-label", "Power");
|
||||
btn.setAttribute("data-testid", "power-button");
|
||||
btn.setAttribute("data-pd-tooltip", "true");
|
||||
btn.setAttribute("title", "Power (Restart / Shutdown)");
|
||||
|
||||
btn.innerHTML = `
|
||||
<div class="side-bar-button-content flex flex-col items-center gap-2">
|
||||
<div class="sidebar-icon-wrapper relative">
|
||||
${getPowerIconSvg(18)}
|
||||
</div>
|
||||
<span class="side-bar-button-label line-clamp-2 w-max max-w-[calc(var(--sidebar-width)-var(--sidebar-padding))] text-center text-2xs wrap-break-word whitespace-normal">Power</span>
|
||||
</div>
|
||||
`;
|
||||
|
||||
btn.onclick = (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
showPowerSidePopup(btn);
|
||||
};
|
||||
|
||||
// Insert at the top of the bottom group (above Help Center)
|
||||
bottomGroup.insertBefore(btn, bottomGroup.firstChild);
|
||||
updateSidebarButtonVisuals();
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: EXTENSION_NAME,
|
||||
async setup() {
|
||||
injectStyles();
|
||||
await fetchPowerStatus();
|
||||
|
||||
// 1. Attach native button in bottom sidebar
|
||||
attachBottomSidebarButton();
|
||||
const observer = new MutationObserver(() => attachBottomSidebarButton());
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
|
||||
// 2. Fallback listener on any matching button
|
||||
setTimeout(() => {
|
||||
const btns = document.querySelectorAll('.comfy-power-btn, [aria-label*="Power"], .side-bar-button:has(.pi-power)');
|
||||
btns.forEach(b => {
|
||||
b.addEventListener("click", (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
showPowerSidePopup(b);
|
||||
}, true);
|
||||
});
|
||||
updateSidebarButtonVisuals();
|
||||
}, 1000);
|
||||
|
||||
try {
|
||||
const handlePowerStatus = (e) => {
|
||||
currentPowerState = e.detail || currentPowerState;
|
||||
updateSidebarButtonVisuals();
|
||||
};
|
||||
api.addEventListener("saturnnodes_power_status", handlePowerStatus);
|
||||
} catch (_) {}
|
||||
|
||||
setInterval(fetchPowerStatus, 3000);
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,314 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.SaturnNodes.PreviewImageSizeAspectRatio",
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== "PreviewImageSizeAspectRatio" && nodeData.name !== "TextAspectRatioFinder") return;
|
||||
|
||||
if (nodeData.name === "PreviewImageSizeAspectRatio") {
|
||||
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function() {
|
||||
if (onNodeCreated) onNodeCreated.apply(this, arguments);
|
||||
|
||||
// STRICTLY REMOVE ALL OUTPUT SLOTS (Once on node creation)
|
||||
this.outputs = [];
|
||||
if (this.outputs) this.outputs.length = 0;
|
||||
|
||||
this.aspect_ratio_data = {
|
||||
ratio: 1.0,
|
||||
display_text: "1 x 1"
|
||||
};
|
||||
|
||||
// Minimum size set to [230, 160] to accommodate 5-digit labels on left and right margins cleanly
|
||||
this.size = [240, 240];
|
||||
this.min_size = [230, 160];
|
||||
this.resizable = true;
|
||||
|
||||
// PERMANENTLY UNLOCK ASPECT RATIO RESIZING
|
||||
Object.defineProperty(this, "aspect_ratio", {
|
||||
get() { return false; },
|
||||
set(v) {},
|
||||
configurable: true
|
||||
});
|
||||
|
||||
// DOM Container for ComfyUI V2 Vue UI
|
||||
const viewContainer = document.createElement("div");
|
||||
viewContainer.className = "saturnnodes-ar-preview-container";
|
||||
viewContainer.style.width = "100%";
|
||||
viewContainer.style.height = "calc(100% - 110px)";
|
||||
viewContainer.style.minHeight = "80px";
|
||||
viewContainer.style.display = "flex";
|
||||
viewContainer.style.alignItems = "center";
|
||||
viewContainer.style.justifyContent = "center";
|
||||
// Equal 65px left and right padding for horizontal balance
|
||||
viewContainer.style.padding = "10px 65px 35px 65px";
|
||||
viewContainer.style.boxSizing = "border-box";
|
||||
viewContainer.style.pointerEvents = "none";
|
||||
viewContainer.style.position = "relative";
|
||||
viewContainer.style.overflow = "hidden";
|
||||
|
||||
viewContainer.innerHTML = `
|
||||
<div class="ar-box-wrapper" style="position: relative; display: flex; align-items: center; justify-content: center; width: 50px; height: 50px; transition: none !important;">
|
||||
<svg class="ar-box-svg" style="position: absolute; top: 0; left: 0; display: block; overflow: visible; pointer-events: none; transition: none !important;" width="50" height="50" viewBox="0 0 50 50">
|
||||
<path class="ar-crop-path" d="" stroke="#ffffff" stroke-width="2.8" stroke-linecap="round" stroke-linejoin="round" fill="none" style="transition: none !important;"/>
|
||||
</svg>
|
||||
<div class="ar-label-height" style="position: absolute; right: 100%; margin-right: 8px; top: 50%; transform: translateY(-50%); font-size: 13px; font-weight: 700; color: #ffffff; font-family: Inter, system-ui, sans-serif; white-space: nowrap;">1</div>
|
||||
<div class="ar-label-width" style="position: absolute; top: 100%; margin-top: 8px; left: 50%; transform: translateX(-50%); font-size: 13px; font-weight: 700; color: #ffffff; font-family: Inter, system-ui, sans-serif; white-space: nowrap;">1</div>
|
||||
</div>
|
||||
`;
|
||||
|
||||
this._last_render_cache = "";
|
||||
|
||||
this.updateDomPreview = function() {
|
||||
if (!viewContainer) return;
|
||||
const data = this.aspect_ratio_data || { ratio: 1.0, display_text: "1 x 1" };
|
||||
const ratio = Math.max(0.05, Math.min(20.0, data.ratio || 1.0));
|
||||
|
||||
let partW = "1";
|
||||
let partH = "1";
|
||||
|
||||
if (data.display_text) {
|
||||
const str = String(data.display_text).trim();
|
||||
if (str.includes(" x ")) {
|
||||
const parts = str.split(" x ");
|
||||
partW = parts[0];
|
||||
partH = parts[1] || "";
|
||||
} else if (str.includes("x")) {
|
||||
const parts = str.split("x");
|
||||
partW = parts[0];
|
||||
partH = parts[1] || "";
|
||||
} else if (str.includes(":")) {
|
||||
const parts = str.split(":");
|
||||
partW = parts[0];
|
||||
partH = parts[1] || "";
|
||||
} else {
|
||||
partW = str;
|
||||
partH = "";
|
||||
}
|
||||
}
|
||||
|
||||
const widthEl = viewContainer.querySelector(".ar-label-width");
|
||||
const heightEl = viewContainer.querySelector(".ar-label-height");
|
||||
const boxWrapper = viewContainer.querySelector(".ar-box-wrapper");
|
||||
const svgEl = viewContainer.querySelector(".ar-box-svg");
|
||||
const pathEl = viewContainer.querySelector(".ar-crop-path");
|
||||
|
||||
if (widthEl) widthEl.innerText = partW;
|
||||
if (heightEl) heightEl.innerText = partH;
|
||||
|
||||
if (boxWrapper && svgEl && pathEl) {
|
||||
const availW = Math.max(20, (this.size[0] || 240) - 130);
|
||||
const availH = Math.max(20, (this.size[1] || 240) - 150);
|
||||
|
||||
let w, h;
|
||||
if (ratio >= availW / availH) {
|
||||
w = availW;
|
||||
h = availW / ratio;
|
||||
} else {
|
||||
h = availH;
|
||||
w = availH * ratio;
|
||||
}
|
||||
|
||||
w = Math.round(w);
|
||||
h = Math.round(h);
|
||||
|
||||
const cacheKey = `${w}_${h}_${partW}_${partH}`;
|
||||
if (this._last_render_cache === cacheKey) return;
|
||||
this._last_render_cache = cacheKey;
|
||||
|
||||
boxWrapper.style.width = `${w}px`;
|
||||
boxWrapper.style.height = `${h}px`;
|
||||
|
||||
svgEl.setAttribute("width", w);
|
||||
svgEl.setAttribute("height", h);
|
||||
svgEl.setAttribute("viewBox", `0 0 ${w} ${h}`);
|
||||
|
||||
const cX = w * 0.20;
|
||||
const cY = h * 0.20;
|
||||
const r = Math.min(12, Math.min(cX, cY) * 0.5);
|
||||
|
||||
const d = `M 0,${cY} L 0,${r} A ${r} ${r} 0 0 1 ${r},0 L ${cX},0 ` +
|
||||
`M ${w - cX},0 L ${w - r},0 A ${r} ${r} 0 0 1 ${w},${r} L ${w},${cY} ` +
|
||||
`M 0,${h - cY} L 0,${h - r} A ${r} ${r} 0 0 0 ${r},${h} L ${cX},${h} ` +
|
||||
`M ${w - cX},${h} L ${w - r},${h} A ${r} ${r} 0 0 0 ${w},${h - r} L ${w},${h - cY} ` +
|
||||
`M ${w * 0.40},0 L ${w * 0.60},0 ` +
|
||||
`M ${w * 0.40},${h} L ${w * 0.60},${h} ` +
|
||||
`M 0,${h * 0.40} L 0,${h * 0.60} ` +
|
||||
`M ${w},${h * 0.40} L ${w},${h * 0.60}`;
|
||||
|
||||
pathEl.setAttribute("d", d);
|
||||
}
|
||||
};
|
||||
|
||||
this.addDOMWidget("ar_preview_display", "HTML", viewContainer, {
|
||||
getValue() { return ""; },
|
||||
setValue(val) {},
|
||||
serialize: false,
|
||||
computeSize() { return [100, 80]; }
|
||||
});
|
||||
|
||||
// Canvas drawing fallback for LiteGraph (Classic V1)
|
||||
const origOnDrawForeground = this.onDrawForeground;
|
||||
this.onDrawForeground = function(ctx) {
|
||||
if (origOnDrawForeground) origOnDrawForeground.apply(this, arguments);
|
||||
if (this.flags?.collapsed) return;
|
||||
|
||||
// HIGH-PERFORMANCE OPTIMIZATION:
|
||||
// If HTML DOM widget is active in document (Vue UI / V2), skip duplicate canvas drawing!
|
||||
if (viewContainer && viewContainer.isConnected && viewContainer.offsetParent !== null) {
|
||||
return;
|
||||
}
|
||||
|
||||
const data = this.aspect_ratio_data || { ratio: 1.0, display_text: "1 x 1" };
|
||||
const ratio = Math.max(0.05, Math.min(20.0, data.ratio || 1.0));
|
||||
|
||||
let partW = "1";
|
||||
let partH = "1";
|
||||
if (data.display_text) {
|
||||
const str = String(data.display_text).trim();
|
||||
if (str.includes(" x ")) {
|
||||
const parts = str.split(" x ");
|
||||
partW = parts[0];
|
||||
partH = parts[1] || "";
|
||||
} else if (str.includes("x")) {
|
||||
const parts = str.split("x");
|
||||
partW = parts[0];
|
||||
partH = parts[1] || "";
|
||||
} else if (str.includes(":")) {
|
||||
const parts = str.split(":");
|
||||
partW = parts[0];
|
||||
partH = parts[1] || "";
|
||||
} else {
|
||||
partW = str;
|
||||
partH = "";
|
||||
}
|
||||
}
|
||||
|
||||
ctx.save();
|
||||
|
||||
const startY = 30 + (this.inputs ? this.inputs.length * 20 : 80);
|
||||
|
||||
const marginLeft = 65;
|
||||
const marginRight = 65;
|
||||
const marginTop = startY + 10;
|
||||
const marginBottom = 35;
|
||||
|
||||
const availW = Math.max(20, this.size[0] - marginLeft - marginRight);
|
||||
const availH = Math.max(20, this.size[1] - marginTop - marginBottom);
|
||||
|
||||
let boxW, boxH;
|
||||
if (ratio >= availW / availH) {
|
||||
boxW = availW;
|
||||
boxH = availW / ratio;
|
||||
} else {
|
||||
boxH = availH;
|
||||
boxW = availH * ratio;
|
||||
}
|
||||
|
||||
const boxX = marginLeft + (availW - boxW) / 2;
|
||||
const boxY = marginTop + (availH - boxH) / 2;
|
||||
|
||||
const cX = boxW * 0.20;
|
||||
const cY = boxH * 0.20;
|
||||
const r = Math.min(12, Math.min(cX, cY) * 0.5);
|
||||
|
||||
ctx.strokeStyle = "#ffffff";
|
||||
ctx.lineWidth = 2.8;
|
||||
ctx.lineCap = "round";
|
||||
ctx.lineJoin = "round";
|
||||
ctx.beginPath();
|
||||
|
||||
// Top-Left Corner
|
||||
ctx.moveTo(boxX, boxY + cY);
|
||||
ctx.lineTo(boxX, boxY + r);
|
||||
ctx.arcTo(boxX, boxY, boxX + r, boxY, r);
|
||||
ctx.lineTo(boxX + cX, boxY);
|
||||
|
||||
// Top-Right Corner
|
||||
ctx.moveTo(boxX + boxW - cX, boxY);
|
||||
ctx.lineTo(boxX + boxW - r, boxY);
|
||||
ctx.arcTo(boxX + boxW, boxY, boxX + boxW, boxY + r, r);
|
||||
ctx.lineTo(boxX + boxW, boxY + cY);
|
||||
|
||||
// Bottom-Left Corner
|
||||
ctx.moveTo(boxX, boxY + boxH - cY);
|
||||
ctx.lineTo(boxX, boxY + boxH - r);
|
||||
ctx.arcTo(boxX, boxY + boxH, boxX + r, boxY + boxH, r);
|
||||
ctx.lineTo(boxX + cX, boxY + boxH);
|
||||
|
||||
// Bottom-Right Corner
|
||||
ctx.moveTo(boxX + boxW - cX, boxY + boxH);
|
||||
ctx.lineTo(boxX + boxW - r, boxY + boxH);
|
||||
ctx.arcTo(boxX + boxW, boxY + boxH, boxX + boxW, boxY + boxH - r, r);
|
||||
ctx.lineTo(boxX + boxW, boxY + boxH - cY);
|
||||
|
||||
// Top Center Handle
|
||||
ctx.moveTo(boxX + boxW * 0.40, boxY);
|
||||
ctx.lineTo(boxX + boxW * 0.60, boxY);
|
||||
|
||||
// Bottom Center Handle
|
||||
ctx.moveTo(boxX + boxW * 0.40, boxY + boxH);
|
||||
ctx.lineTo(boxX + boxW * 0.60, boxY + boxH);
|
||||
|
||||
// Left Center Handle
|
||||
ctx.moveTo(boxX, boxY + boxH * 0.40);
|
||||
ctx.lineTo(boxX, boxY + boxH * 0.60);
|
||||
|
||||
// Right Center Handle
|
||||
ctx.moveTo(boxX + boxW, boxY + boxH * 0.40);
|
||||
ctx.lineTo(boxX + boxW, boxY + boxH * 0.60);
|
||||
|
||||
ctx.stroke();
|
||||
|
||||
ctx.fillStyle = "#ffffff";
|
||||
ctx.font = "bold 13px Inter, -apple-system, sans-serif";
|
||||
|
||||
if (partH) {
|
||||
ctx.textAlign = "right";
|
||||
ctx.textBaseline = "middle";
|
||||
ctx.fillText(partH, boxX - 8, boxY + boxH / 2);
|
||||
}
|
||||
|
||||
if (partW) {
|
||||
ctx.textAlign = "center";
|
||||
ctx.textBaseline = "top";
|
||||
ctx.fillText(partW, boxX + boxW / 2, boxY + boxH + 8);
|
||||
}
|
||||
|
||||
ctx.restore();
|
||||
};
|
||||
|
||||
nodeType.prototype.computeSize = function() {
|
||||
return [230, 160];
|
||||
};
|
||||
|
||||
nodeType.prototype.onResize = function(size) {
|
||||
this.size[0] = Math.max(230, size[0]);
|
||||
this.size[1] = Math.max(160, size[1]);
|
||||
if (typeof this.updateDomPreview === "function") {
|
||||
this.updateDomPreview();
|
||||
}
|
||||
};
|
||||
};
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
api.addEventListener("saturnnodes_update_preview_aspect_ratio", (event) => {
|
||||
const detail = event.detail;
|
||||
if (!detail || !detail.node_id) return;
|
||||
|
||||
const node = app.graph.getNodeById(detail.node_id);
|
||||
if (node && (node.type === "PreviewImageSizeAspectRatio" || node.type === "TextAspectRatioFinder")) {
|
||||
node.aspect_ratio_data = {
|
||||
ratio: detail.ratio,
|
||||
display_text: detail.display_text
|
||||
};
|
||||
if (typeof node.updateDomPreview === "function") {
|
||||
node.updateDomPreview();
|
||||
}
|
||||
app.graph.setDirtyCanvas(true, true);
|
||||
}
|
||||
});
|
||||
@@ -66,8 +66,8 @@ class PreviewManagerClass {
|
||||
const images = document.querySelectorAll('img[alt*="Live sampling preview"]');
|
||||
images.forEach(imgEl => {
|
||||
const src = imgEl.src;
|
||||
if (src && src.startsWith("blob:") && imgEl.dataset.lastFlowControlSrc !== src) {
|
||||
imgEl.dataset.lastFlowControlSrc = src;
|
||||
if (src && src.startsWith("blob:") && imgEl.dataset.lastSaturnNodesSrc !== src) {
|
||||
imgEl.dataset.lastSaturnNodesSrc = src;
|
||||
|
||||
const nodeEl = imgEl.closest('[data-node-id], [data-widgets-grid-node-id], .lg-node');
|
||||
const nodeId = nodeEl ? (nodeEl.dataset.nodeId || nodeEl.dataset.widgetsGridNodeId || nodeEl.getAttribute('data-node-id')) : null;
|
||||
|
||||
@@ -0,0 +1,302 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { isSaturnColorsEnabled } from "../saturnnodes_colors.js";
|
||||
import { resolveLiveWidgetElement } from "./live_dom.js";
|
||||
|
||||
function countPromptsInText(text, separator, customRegex) {
|
||||
if (!text || !text.trim()) return 0;
|
||||
const cleanText = text.replace(/\r\n/g, '\n').replace(/\r/g, '\n');
|
||||
const sep = (separator || ">1 Empty Line").toLowerCase();
|
||||
|
||||
if (sep.includes("regex")) {
|
||||
const patStr = (customRegex || "").trim() || "\\n\\s*\\n+";
|
||||
try {
|
||||
const re = new RegExp(patStr, "g");
|
||||
const rawBlocks = cleanText.split(re);
|
||||
let count = 0;
|
||||
for (const b of rawBlocks) {
|
||||
const s = b.trim();
|
||||
if (s) {
|
||||
try {
|
||||
if (!new RegExp(`^(?:${patStr})$`).test(b)) {
|
||||
count++;
|
||||
}
|
||||
} catch (_) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
}
|
||||
return count;
|
||||
} catch (e) {
|
||||
// Fallback on invalid regex
|
||||
const rawBlocks = cleanText.split(/\n\s*\n+/);
|
||||
return rawBlocks.filter(b => b.trim()).length;
|
||||
}
|
||||
} else if (sep.includes("newline")) {
|
||||
return cleanText.split('\n').filter(line => line.trim()).length;
|
||||
} else if (sep.includes(">2") || sep.includes("2 empty")) {
|
||||
return cleanText.split(/(?:\n\s*){3,}/).filter(b => b.trim()).length;
|
||||
} else { // >1 Empty Line / >1 Emptyline
|
||||
return cleanText.split(/\n\s*\n+/).filter(b => b.trim()).length;
|
||||
}
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.SaturnNodes.PromptCounter",
|
||||
async nodeCreated(node) {
|
||||
if (node.comfyClass === "PromptCounter") {
|
||||
const textWidget = node.widgets?.find(w => w.name === "text");
|
||||
const sepWidget = node.widgets?.find(w => w.name === "separator");
|
||||
const regexWidget = node.widgets?.find(w => w.name === "custom_regex");
|
||||
|
||||
// Live Prompt Count preview widget matching ResolutionSelector styling
|
||||
let countWidget = node.widgets?.find(w => w.name === "prompt_count_preview" || w._isPromptCount);
|
||||
let badgeSpan = node._promptBadgeSpan;
|
||||
if (!countWidget) {
|
||||
const badgeEl = document.createElement("div");
|
||||
badgeEl.className = "not-disabled:bg-component-node-widget-background not-disabled:text-component-node-foreground [[readonly]]:bg-component-node-widget-background-disabled border-none rounded-md flex w-full items-center justify-center gap-2 px-2 h-6 col-span-2";
|
||||
badgeEl.setAttribute("data-widget-name", "prompt_count_preview");
|
||||
badgeEl.setAttribute("node-type", "PromptCounter");
|
||||
badgeEl.style.cssText = "display: flex; align-items: center; justify-content: center; gap: 6px; padding: 0 8px; height: 24px !important; min-height: 24px !important; max-height: 24px !important; flex-grow: 0 !important; flex-shrink: 0 !important; align-self: center !important; overflow: hidden !important; border-radius: 6px; background: var(--component-node-widget-background, rgba(255, 255, 255, 0.06)); color: var(--component-node-foreground, #e4e4e7); font-size: 12px; font-weight: 500; width: 100%; box-sizing: border-box; user-select: none;";
|
||||
|
||||
badgeSpan = document.createElement("span");
|
||||
badgeSpan.className = "text-xs";
|
||||
badgeSpan.style.cssText = "font-weight: 600;";
|
||||
badgeSpan.textContent = "0 Prompts";
|
||||
badgeEl.appendChild(badgeSpan);
|
||||
node._promptBadgeSpan = badgeSpan;
|
||||
node._promptBadgeEl = badgeEl;
|
||||
|
||||
const enforceFixedBadgeSize = () => {
|
||||
badgeEl.style.setProperty("height", "24px", "important");
|
||||
badgeEl.style.setProperty("min-height", "24px", "important");
|
||||
badgeEl.style.setProperty("max-height", "24px", "important");
|
||||
badgeEl.style.setProperty("flex-grow", "0", "important");
|
||||
badgeEl.style.setProperty("flex-shrink", "0", "important");
|
||||
badgeEl.style.setProperty("align-self", "center", "important");
|
||||
const widgetContainer = badgeEl.closest?.(".lg-node-widget");
|
||||
if (widgetContainer) {
|
||||
widgetContainer.style.setProperty("align-self", "center", "important");
|
||||
widgetContainer.style.setProperty("flex-grow", "0", "important");
|
||||
}
|
||||
};
|
||||
node._enforceFixedBadgeSize = enforceFixedBadgeSize;
|
||||
|
||||
countWidget = node.addDOMWidget("prompt_count_preview", "preview", badgeEl, {
|
||||
serialize: false,
|
||||
getValue() { return badgeSpan.textContent; },
|
||||
setValue(v) { badgeSpan.textContent = v; }
|
||||
});
|
||||
countWidget._isPromptCount = true;
|
||||
}
|
||||
|
||||
if (countWidget) {
|
||||
// Crucial: shadow computeLayoutSize so ComfyUI Vue deriveWidgetRenderState evaluates hasLayoutSize to false.
|
||||
// This forces grid-template-rows to use 'min-content' (24px) instead of 'auto' (which was stealing half the node height).
|
||||
Object.defineProperty(countWidget, "computeLayoutSize", {
|
||||
value: undefined,
|
||||
writable: true,
|
||||
configurable: true
|
||||
});
|
||||
countWidget.hasLayoutSize = false;
|
||||
countWidget.computeSize = function() {
|
||||
return [node.size ? (node.size[0] - 20) : 200, 24];
|
||||
};
|
||||
}
|
||||
|
||||
function updatePromptBadgeColor() {
|
||||
const enabled = isSaturnColorsEnabled();
|
||||
const badgeEl = resolveLiveWidgetElement(
|
||||
node._promptBadgeEl,
|
||||
node.id,
|
||||
"prompt_count_preview"
|
||||
);
|
||||
if (badgeEl) node._promptBadgeEl = badgeEl;
|
||||
const badgeSpan = badgeEl?.querySelector?.("span") || node._promptBadgeSpan;
|
||||
if (badgeSpan) node._promptBadgeSpan = badgeSpan;
|
||||
|
||||
if (badgeSpan) {
|
||||
badgeSpan.style.color = enabled ? "#f59e0b" : "var(--text-muted, #a1a1aa)";
|
||||
}
|
||||
if (badgeEl) {
|
||||
badgeEl.style.background = enabled
|
||||
? "rgba(180, 140, 95, 0.12)"
|
||||
: "var(--component-node-widget-background, rgba(255, 255, 255, 0.06))";
|
||||
badgeEl.style.border = enabled
|
||||
? "1px solid rgba(180, 140, 95, 0.35)"
|
||||
: "none";
|
||||
}
|
||||
}
|
||||
node._updateCounterBadgeColor = updatePromptBadgeColor;
|
||||
|
||||
function updateCount() {
|
||||
const text = (textWidget?.inputEl ? textWidget.inputEl.value : textWidget?.value) || "";
|
||||
const sep = sepWidget?.value || ">1 Empty Line";
|
||||
const regex = regexWidget?.value || "";
|
||||
const count = countPromptsInText(text, sep, regex);
|
||||
node._promptCount = count;
|
||||
const label = `${count} Prompt${count === 1 ? "" : "s"}`;
|
||||
const liveBadgeEl = resolveLiveWidgetElement(
|
||||
node._promptBadgeEl,
|
||||
node.id,
|
||||
"prompt_count_preview"
|
||||
);
|
||||
if (liveBadgeEl) node._promptBadgeEl = liveBadgeEl;
|
||||
const liveBadgeSpan = liveBadgeEl?.querySelector?.("span") || node._promptBadgeSpan;
|
||||
if (liveBadgeSpan) {
|
||||
node._promptBadgeSpan = liveBadgeSpan;
|
||||
liveBadgeSpan.textContent = label;
|
||||
}
|
||||
if (countWidget) {
|
||||
countWidget.value = label;
|
||||
}
|
||||
updatePromptBadgeColor();
|
||||
|
||||
// Manage custom_regex enable/disable
|
||||
if (regexWidget) {
|
||||
const isCustom = String(sep || "").toLowerCase().includes("regex");
|
||||
regexWidget.disabled = !isCustom;
|
||||
if (regexWidget.inputEl) {
|
||||
regexWidget.inputEl.disabled = !isCustom;
|
||||
regexWidget.inputEl.style.opacity = isCustom ? "1" : "0.4";
|
||||
regexWidget.inputEl.style.pointerEvents = isCustom ? "auto" : "none";
|
||||
}
|
||||
if (regexWidget.options) {
|
||||
regexWidget.options.disabled = !isCustom;
|
||||
}
|
||||
}
|
||||
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
}
|
||||
|
||||
// Hook DOM input listener on text textarea
|
||||
function hookInputEl() {
|
||||
if (textWidget?.inputEl && !textWidget.inputEl._promptCounterHooked) {
|
||||
textWidget.inputEl._promptCounterHooked = true;
|
||||
textWidget.inputEl.addEventListener("input", () => {
|
||||
textWidget.value = textWidget.inputEl.value;
|
||||
updateCount();
|
||||
});
|
||||
textWidget.inputEl.addEventListener("paste", () => {
|
||||
setTimeout(() => {
|
||||
textWidget.value = textWidget.inputEl.value;
|
||||
updateCount();
|
||||
}, 20);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (textWidget) {
|
||||
const origTextCb = textWidget.callback;
|
||||
textWidget.callback = function(val) {
|
||||
if (origTextCb) origTextCb.apply(this, arguments);
|
||||
updateCount();
|
||||
};
|
||||
}
|
||||
|
||||
if (sepWidget) {
|
||||
const origSepCb = sepWidget.callback;
|
||||
sepWidget.callback = function(val) {
|
||||
if (origSepCb) origSepCb.apply(this, arguments);
|
||||
updateCount();
|
||||
};
|
||||
}
|
||||
|
||||
if (regexWidget) {
|
||||
const origRegexCb = regexWidget.callback;
|
||||
regexWidget.callback = function(val) {
|
||||
if (origRegexCb) origRegexCb.apply(this, arguments);
|
||||
updateCount();
|
||||
};
|
||||
|
||||
const origMouse = regexWidget.mouse;
|
||||
regexWidget.mouse = function(event, pos, n) {
|
||||
if (this.disabled) return false;
|
||||
if (origMouse) return origMouse.apply(this, arguments);
|
||||
};
|
||||
|
||||
const origDraw = regexWidget.draw;
|
||||
regexWidget.draw = function(ctx, n, widget_width, y, widget_height) {
|
||||
if (this.disabled) {
|
||||
ctx.save();
|
||||
ctx.globalAlpha = 0.35;
|
||||
if (origDraw) origDraw.apply(this, arguments);
|
||||
ctx.restore();
|
||||
} else {
|
||||
if (origDraw) origDraw.apply(this, arguments);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// Also draw title pill badge when node is collapsed
|
||||
const origOnDrawForeground = node.onDrawForeground;
|
||||
node.onDrawForeground = function(ctx) {
|
||||
hookInputEl();
|
||||
|
||||
// If textWidget value changed without callback
|
||||
const currentVal = (textWidget?.inputEl ? textWidget.inputEl.value : textWidget?.value) || "";
|
||||
if (currentVal !== node._lastPromptVal) {
|
||||
node._lastPromptVal = currentVal;
|
||||
updateCount();
|
||||
}
|
||||
|
||||
if (node.flags?.collapsed) {
|
||||
const isSaturn = isSaturnColorsEnabled();
|
||||
const count = node._promptCount ?? 0;
|
||||
const label = `${count} Prompt${count === 1 ? "" : "s"}`;
|
||||
ctx.save();
|
||||
ctx.font = "bold 11px Inter, system-ui, sans-serif";
|
||||
const badgeText = label;
|
||||
const tw = ctx.measureText(badgeText).width;
|
||||
const bx = node.size[0] - tw - 16;
|
||||
const by = -LiteGraph.NODE_TITLE_HEIGHT + 3;
|
||||
|
||||
if (isSaturn) {
|
||||
ctx.fillStyle = count > 0 ? "rgba(180, 140, 95, 0.25)" : "rgba(255, 255, 255, 0.08)";
|
||||
ctx.strokeStyle = count > 0 ? "rgba(180, 140, 95, 0.6)" : "rgba(255, 255, 255, 0.2)";
|
||||
} else {
|
||||
ctx.fillStyle = count > 0 ? "rgba(255, 255, 255, 0.12)" : "rgba(255, 255, 255, 0.08)";
|
||||
ctx.strokeStyle = count > 0 ? "rgba(255, 255, 255, 0.25)" : "rgba(255, 255, 255, 0.2)";
|
||||
}
|
||||
ctx.lineWidth = 1;
|
||||
ctx.beginPath();
|
||||
if (ctx.roundRect) {
|
||||
ctx.roundRect(bx, by, tw + 10, 18, 4);
|
||||
} else {
|
||||
ctx.rect(bx, by, tw + 10, 18);
|
||||
}
|
||||
ctx.fill();
|
||||
ctx.stroke();
|
||||
|
||||
ctx.fillStyle = isSaturn ? (count > 0 ? "#decbb2" : "#cbd5e1") : (count > 0 ? "#e2e8f0" : "#cbd5e1");
|
||||
ctx.textAlign = "left";
|
||||
ctx.textBaseline = "middle";
|
||||
ctx.fillText(badgeText, bx + 5, by + 9);
|
||||
ctx.restore();
|
||||
}
|
||||
|
||||
if (origOnDrawForeground) origOnDrawForeground.apply(this, arguments);
|
||||
};
|
||||
|
||||
const origOnResize = node.onResize;
|
||||
node.onResize = function() {
|
||||
if (origOnResize) origOnResize.apply(this, arguments);
|
||||
node._enforceFixedBadgeSize?.();
|
||||
};
|
||||
|
||||
const origOnConfigure = node.onConfigure;
|
||||
node.onConfigure = function() {
|
||||
if (origOnConfigure) origOnConfigure.apply(this, arguments);
|
||||
setTimeout(() => {
|
||||
hookInputEl();
|
||||
node._enforceFixedBadgeSize?.();
|
||||
updateCount();
|
||||
}, 40);
|
||||
};
|
||||
|
||||
setTimeout(() => {
|
||||
hookInputEl();
|
||||
updateCount();
|
||||
}, 20);
|
||||
}
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,172 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
import { authenticatedFetch } from "./auth_helper.js";
|
||||
|
||||
function findGraphNode(id) {
|
||||
if (!id && id !== 0) return null;
|
||||
const strId = String(id);
|
||||
const numId = Number(id);
|
||||
|
||||
let node = app.graph?.getNodeById?.(id) || app.graph?.getNodeById?.(numId) || app.graph?.getNodeById?.(strId);
|
||||
if (node) return node;
|
||||
|
||||
if (app.graph && Array.isArray(app.graph._nodes)) {
|
||||
for (const n of app.graph._nodes) {
|
||||
if (String(n.id) === strId) return n;
|
||||
if (n.subgraph && typeof n.subgraph.getNodeById === "function") {
|
||||
const sub = n.subgraph.getNodeById(id) || n.subgraph.getNodeById(numId) || n.subgraph.getNodeById(strId);
|
||||
if (sub) return sub;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (app.graph && app.graph._subgraphs) {
|
||||
const subgraphs = Array.isArray(app.graph._subgraphs)
|
||||
? app.graph._subgraphs
|
||||
: Object.values(app.graph._subgraphs);
|
||||
for (const sg of subgraphs) {
|
||||
if (sg && typeof sg.getNodeById === "function") {
|
||||
const sub = sg.getNodeById(id) || sg.getNodeById(numId) || sg.getNodeById(strId);
|
||||
if (sub) return sub;
|
||||
}
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.SaturnNodes.PromptQueueIterator",
|
||||
async setup() {
|
||||
if (api && api.addEventListener) {
|
||||
const handleProgress = (event) => {
|
||||
const data = event.detail;
|
||||
if (!data || !data.node_id) return;
|
||||
const node = findGraphNode(data.node_id);
|
||||
if (node && node.widgets) {
|
||||
const statusWidget = node.widgets.find(w => w.name === "Progress");
|
||||
if (statusWidget) {
|
||||
statusWidget.value = `Run ${data.current_run} / ${data.total_runs}`;
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
}
|
||||
}
|
||||
};
|
||||
api.addEventListener("saturnnodes_prompt_iterator_progress", handleProgress);
|
||||
}
|
||||
},
|
||||
async nodeCreated(node) {
|
||||
if (node.comfyClass === "PromptQueueIterator") {
|
||||
// 1. Live Progress Status Widget
|
||||
const statusWidget = node.addWidget("text", "Progress", "Ready (-- / --)", () => {}, { serialize: false });
|
||||
statusWidget.disabled = true;
|
||||
|
||||
// 2. Reset Counter (0) Button
|
||||
const resetBtn = node.addWidget("button", "Reset Counter (0)", null, async () => {
|
||||
const proceed = confirm(
|
||||
"Reset this node's prompt counter to 0?\n\n" +
|
||||
"• New queued runs start again from the FIRST prompt.\n" +
|
||||
"• Already-assigned running/queued items keep their prompt (crash recovery).\n" +
|
||||
"• The state file is NOT emptied - use Settings → \"Reset All Queues\" for a full wipe.\n\nContinue?"
|
||||
);
|
||||
if (!proceed) return;
|
||||
try {
|
||||
resetBtn.name = "⏳ Resetting...";
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
const resp = await authenticatedFetch("/saturnnodes/prompt_iterator/reset_node", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ node_id: String(node.id) })
|
||||
});
|
||||
const res = await resp.json();
|
||||
if (res && res.status === "ok") {
|
||||
resetBtn.name = "✅ Reset to 0!";
|
||||
statusWidget.value = "Reset (0 / --)";
|
||||
} else {
|
||||
resetBtn.name = "⚠️ Error";
|
||||
}
|
||||
} catch (e) {
|
||||
console.error("[SaturnNodes] Error resetting counter:", e);
|
||||
resetBtn.name = "❌ Failed";
|
||||
}
|
||||
setTimeout(() => {
|
||||
resetBtn.name = "Reset Counter (0)";
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
}, 1200);
|
||||
});
|
||||
resetBtn.serialize = false;
|
||||
|
||||
// 3. Open Queue File in Default OS Editor Button
|
||||
const openFileBtn = node.addWidget("button", "Open Queue File", null, async () => {
|
||||
try {
|
||||
openFileBtn.name = "⏳ Opening...";
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
const resp = await authenticatedFetch("/saturnnodes/prompt_iterator/open_file", { method: "POST" });
|
||||
const res = await resp.json();
|
||||
if (res && res.status === "ok") {
|
||||
openFileBtn.name = "✅ Opened in Editor!";
|
||||
} else {
|
||||
openFileBtn.name = "⚠️ Error";
|
||||
}
|
||||
} catch (e) {
|
||||
console.error("[SaturnNodes] Error opening state file:", e);
|
||||
openFileBtn.name = "❌ Failed";
|
||||
}
|
||||
setTimeout(() => {
|
||||
openFileBtn.name = "Open Queue File";
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
}, 1500);
|
||||
});
|
||||
// 4. Manage custom_regex enable/disable based on separator
|
||||
const sepWidget = node.widgets?.find(w => w.name === "separator");
|
||||
const regexWidget = node.widgets?.find(w => w.name === "custom_regex");
|
||||
|
||||
function updateRegexState() {
|
||||
if (!sepWidget || !regexWidget) return;
|
||||
const isCustom = String(sepWidget.value || "").toLowerCase().includes("regex");
|
||||
regexWidget.disabled = !isCustom;
|
||||
if (regexWidget.inputEl) {
|
||||
regexWidget.inputEl.disabled = !isCustom;
|
||||
regexWidget.inputEl.style.opacity = isCustom ? "1" : "0.4";
|
||||
regexWidget.inputEl.style.pointerEvents = isCustom ? "auto" : "none";
|
||||
}
|
||||
if (regexWidget.options) {
|
||||
regexWidget.options.disabled = !isCustom;
|
||||
}
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
}
|
||||
|
||||
if (sepWidget && regexWidget) {
|
||||
const origSepCb = sepWidget.callback;
|
||||
sepWidget.callback = function(val) {
|
||||
if (origSepCb) origSepCb.apply(this, arguments);
|
||||
updateRegexState();
|
||||
};
|
||||
|
||||
const origMouse = regexWidget.mouse;
|
||||
regexWidget.mouse = function(event, pos, node) {
|
||||
if (this.disabled) return false;
|
||||
if (origMouse) return origMouse.apply(this, arguments);
|
||||
};
|
||||
|
||||
const origDraw = regexWidget.draw;
|
||||
regexWidget.draw = function(ctx, node, widget_width, y, widget_height) {
|
||||
if (this.disabled) {
|
||||
ctx.save();
|
||||
ctx.globalAlpha = 0.35;
|
||||
if (origDraw) origDraw.apply(this, arguments);
|
||||
ctx.restore();
|
||||
} else {
|
||||
if (origDraw) origDraw.apply(this, arguments);
|
||||
}
|
||||
};
|
||||
|
||||
setTimeout(updateRegexState, 20);
|
||||
}
|
||||
|
||||
const origOnConfigure = node.onConfigure;
|
||||
node.onConfigure = function() {
|
||||
if (origOnConfigure) origOnConfigure.apply(this, arguments);
|
||||
setTimeout(updateRegexState, 40);
|
||||
};
|
||||
}
|
||||
}
|
||||
});
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user