Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b21497f664 | ||
|
|
29a8d9286f | ||
|
|
d3ccc24055 | ||
|
|
8756d720ae | ||
|
|
33e9201ac8 | ||
|
|
3f2a1d5421 | ||
|
|
49f8cfbde6 | ||
|
|
913fc5354d | ||
|
|
fb8e8ec8d5 | ||
|
|
59b41a19a4 | ||
|
|
aaf49230b2 |
+10
-3
@@ -9,18 +9,25 @@ CIVITAI_API_KEY=
|
||||
TMDB_API_KEY=
|
||||
|
||||
# Feature toggles
|
||||
ENABLE_CIVITAI_SCRAPING=true
|
||||
ENABLE_TMDB_SCRAPING=true
|
||||
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=Paused
|
||||
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
|
||||
|
||||
@@ -51,6 +51,7 @@ ComfyUI-SaturnNodes/
|
||||
│ ├── 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
|
||||
@@ -238,6 +239,7 @@ STATE_FILE = os.path.join(USER_DIR, "your_state.json")
|
||||
- `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
|
||||
|
||||
+19
-3
@@ -5,29 +5,45 @@ 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.1] - In Development
|
||||
## [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.1`.
|
||||
- **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.
|
||||
- **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 (119 unit tests passing).
|
||||
- **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.
|
||||
|
||||
@@ -107,6 +107,8 @@ Visual thumbnail browser for image folders with instant preview selection, indep
|
||||
- **`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.
|
||||
@@ -242,6 +244,7 @@ Deterministically iterates over multiline prompt text blocks per queue run with
|
||||
|
||||
#### 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.
|
||||
@@ -249,8 +252,11 @@ Deterministically iterates over multiline prompt text blocks per queue run with
|
||||
|
||||
#### Outputs
|
||||
- **`prompt`** (`STRING`): The selected prompt block for the current queue run.
|
||||
- **`remaining_text`** (`STRING`): All remaining prompt blocks joined by the delimiter.
|
||||
- **`remaining_count`** (`INT`): Count of remaining items in queue.
|
||||
- **`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>
|
||||
@@ -374,12 +380,13 @@ Configure options directly under ComfyUI Settings (⚙ gear icon):
|
||||
- **`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: true*): Automatically download model preview images from Civitai via SHA256 file hashes. (Note: Local SHA256 hash searching always works).
|
||||
- **`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>
|
||||
@@ -411,6 +418,7 @@ Configure options directly under ComfyUI Settings (⚙ gear icon):
|
||||
|
||||
- **`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.
|
||||
@@ -424,6 +432,7 @@ Configure options directly under ComfyUI Settings (⚙ gear icon):
|
||||
|
||||
- **`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>
|
||||
@@ -431,6 +440,9 @@ Configure options directly under ComfyUI Settings (⚙ gear icon):
|
||||
|
||||
- **`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>
|
||||
|
||||
+11
-4
@@ -41,14 +41,18 @@ SaturnNodes automatically tracks which workflows are submitted together as a bat
|
||||
|
||||
### 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 (in a paused state, so it doesn't overwhelm your GPU before you're ready).
|
||||
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!
|
||||
- You can adjust the restored image count or toggle this feature in **ComfyUI Settings (⚙️)** under **`Restore Assets on Launch`** and **`Restored Assets Count`** (search for `SaturnNodes`).
|
||||
- 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:
|
||||
@@ -94,7 +98,9 @@ Managing multi-prompt queues and setting the exact ComfyUI batch size has never
|
||||
- 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.
|
||||
- Includes live progress counter widget (`Run X / Y`), an index reset button (`Reset Counter (0)`, now with a confirmation explaining that it only restarts the run pointer — 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`).
|
||||
- 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`).
|
||||
|
||||
---
|
||||
|
||||
@@ -115,7 +121,8 @@ Several nodes across this suite (e.g. *Load Image From Folder*, *Recent Outputs*
|
||||
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).
|
||||
- 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!
|
||||
|
||||
+98
-41
@@ -4,10 +4,11 @@ import subprocess
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
|
||||
__version__ = "2.6.1"
|
||||
__version__ = "2.6.2"
|
||||
|
||||
from .nodes.queue_control import (
|
||||
setup_queue_control_routes,
|
||||
pause_manager,
|
||||
tray_manager,
|
||||
persistent_queue_manager,
|
||||
is_persistent_queue_enabled
|
||||
@@ -290,49 +291,42 @@ async def get_image_prompt_endpoint(request):
|
||||
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)
|
||||
civitai_key = os.getenv("CIVITAI_API_KEY", "")
|
||||
tmdb_key = os.getenv("TMDB_API_KEY", "")
|
||||
enable_tray = os.getenv("ENABLE_TRAY_ICON", "false").lower() in ["true", "1", "yes"]
|
||||
enable_assets_restore = os.getenv("ENABLE_ASSETS_RESTORE", "true").lower() in ["true", "1", "yes"]
|
||||
restore_assets_count = int(os.getenv("RESTORE_ASSETS_COUNT", "64"))
|
||||
enable_process_management = os.getenv("ALLOW_PROCESS_MANAGEMENT", "false").lower() in ["true", "1", "yes"]
|
||||
enable_local_file_execution = os.getenv("ENABLE_LOCAL_FILE_EXECUTION", "false").lower() in ["true", "1", "yes"]
|
||||
if os.path.exists(ENV_FILE):
|
||||
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()
|
||||
elif line.startswith("ENABLE_TRAY_ICON="):
|
||||
enable_tray = line.split("=", 1)[1].strip().lower() in ["true", "1", "yes"]
|
||||
elif line.startswith("ENABLE_ASSETS_RESTORE="):
|
||||
enable_assets_restore = line.split("=", 1)[1].strip().lower() in ["true", "1", "yes"]
|
||||
elif line.startswith("RESTORE_ASSETS_COUNT="):
|
||||
try:
|
||||
restore_assets_count = int(line.split("=", 1)[1].strip())
|
||||
except Exception:
|
||||
pass
|
||||
elif line.startswith("ALLOW_PROCESS_MANAGEMENT="):
|
||||
enable_process_management = line.split("=", 1)[1].strip().lower() in ["true", "1", "yes"]
|
||||
elif line.startswith("ENABLE_LOCAL_FILE_EXECUTION="):
|
||||
enable_local_file_execution = line.split("=", 1)[1].strip().lower() in ["true", "1", "yes"]
|
||||
except Exception:
|
||||
pass
|
||||
# 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,
|
||||
"enable_tray_icon": enable_tray,
|
||||
"enable_assets_restore": enable_assets_restore,
|
||||
"restore_assets_count": restore_assets_count,
|
||||
"enable_process_management": enable_process_management,
|
||||
"enable_local_file_execution": enable_local_file_execution
|
||||
"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)
|
||||
})
|
||||
|
||||
def _clean_env_val(v):
|
||||
@@ -361,6 +355,11 @@ async def save_settings(request):
|
||||
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):
|
||||
@@ -386,6 +385,11 @@ async def save_settings(request):
|
||||
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:
|
||||
@@ -433,6 +437,21 @@ async def save_settings(request):
|
||||
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
|
||||
@@ -469,6 +488,16 @@ async def save_settings(request):
|
||||
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")
|
||||
|
||||
temp_env = ENV_FILE + ".tmp"
|
||||
with open(temp_env, "w", encoding="utf-8") as f:
|
||||
@@ -497,6 +526,29 @@ async def save_settings(request):
|
||||
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"[SaturnNodes] Error saving settings to .env: {e}")
|
||||
@@ -604,14 +656,19 @@ async def export_debug_profile(request):
|
||||
"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", "true"),
|
||||
"enable_tmdb_scraping": get_env_setting("ENABLE_TMDB_SCRAPING", "true"),
|
||||
"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()),
|
||||
|
||||
+294
-153
@@ -1,29 +1,14 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
import hashlib
|
||||
import heapq
|
||||
import datetime
|
||||
import logging
|
||||
from PIL import Image
|
||||
|
||||
import folder_paths
|
||||
from .utils import get_saturnnodes_user_dir
|
||||
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
ENV_FILE = os.path.join(USER_DIR, ".env")
|
||||
|
||||
def get_env_setting(key, default_val):
|
||||
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(f"{key}="):
|
||||
return line.strip().split("=", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
return default_val
|
||||
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()
|
||||
@@ -36,83 +21,222 @@ def get_assets_restore_count():
|
||||
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 _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 body.decode("utf-16-be", errors="ignore").strip("\x00").strip()
|
||||
except Exception:
|
||||
pass
|
||||
if prefix.startswith(b"ASCII"):
|
||||
try:
|
||||
return body.decode("ascii", errors="ignore").strip("\x00").strip()
|
||||
except Exception:
|
||||
pass
|
||||
if prefix.startswith(b"JIS"):
|
||||
try:
|
||||
return body.decode("shift_jis", errors="ignore").strip("\x00").strip()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
return raw.decode("utf-8", errors="ignore").strip("\x00").strip()
|
||||
except Exception:
|
||||
return ""
|
||||
return str(raw).strip()
|
||||
|
||||
def parse_a1111_positive_prompt(parameters):
|
||||
"""Extract the positive prompt from an A1111/Forge `parameters` string."""
|
||||
if not parameters:
|
||||
return ""
|
||||
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.last_debug_report = {}
|
||||
self.restored_count = 0
|
||||
|
||||
def extract_image_metadata(self, filepath):
|
||||
"""Extracts embedded prompt and workflow graph metadata from image (read-only)."""
|
||||
"""Extracts embedded prompt/workflow metadata from PNG, JPEG or WebP."""
|
||||
prompt_dict = {}
|
||||
workflow_dict = {}
|
||||
positive_prompt = ""
|
||||
|
||||
try:
|
||||
with Image.open(filepath) as img:
|
||||
if img.format == "PNG":
|
||||
prompt_raw = img.info.get("prompt")
|
||||
workflow_raw = img.info.get("workflow")
|
||||
parameters_raw = img.info.get("parameters")
|
||||
prompt_raw = img.info.get("prompt")
|
||||
workflow_raw = img.info.get("workflow")
|
||||
parameters_raw = img.info.get("parameters")
|
||||
|
||||
if prompt_raw:
|
||||
# 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:
|
||||
prompt_dict = json.loads(prompt_raw) if isinstance(prompt_raw, str) else prompt_raw
|
||||
parameters_raw = _decode_exif_text(exif.get(0x9286))
|
||||
except Exception:
|
||||
pass
|
||||
parameters_raw = ""
|
||||
|
||||
if workflow_raw:
|
||||
try:
|
||||
workflow_dict = json.loads(workflow_raw) if isinstance(workflow_raw, str) else workflow_raw
|
||||
except Exception:
|
||||
pass
|
||||
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 parameters_raw:
|
||||
positive_prompt = str(parameters_raw)
|
||||
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):
|
||||
"""Scans output directory recursively (strictly read-only) and returns top `limit` newest images."""
|
||||
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", ".PNG", ".JPG", ".JPEG", ".WEBP"}
|
||||
found_images = []
|
||||
|
||||
valid_extensions = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
norm_base = os.path.normpath(base_output_dir)
|
||||
if stats is None:
|
||||
stats = {}
|
||||
|
||||
for root, _, files in os.walk(norm_base):
|
||||
for file in files:
|
||||
_, ext = os.path.splitext(file)
|
||||
if ext in valid_extensions:
|
||||
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)
|
||||
rel_dir = os.path.relpath(root, norm_base)
|
||||
subfolder = "" if rel_dir == "." else rel_dir.replace("\\", "/")
|
||||
found_images.append({
|
||||
"filepath": full_path,
|
||||
"filename": file,
|
||||
"subfolder": subfolder,
|
||||
"mtime": mtime
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
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,
|
||||
}
|
||||
|
||||
# Sort newest first, take limit, then reverse to chronological (oldest to newest)
|
||||
found_images.sort(key=lambda x: x["mtime"], reverse=True)
|
||||
top_images = found_images[:limit]
|
||||
# Keep only the top-N candidates (bounded memory), then restore
|
||||
# chronological order (oldest -> newest) for stable history numbering.
|
||||
top_images = heapq.nlargest(max(1, int(limit)), iter_images(), key=lambda x: x["mtime"])
|
||||
top_images.reverse()
|
||||
|
||||
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):
|
||||
"""Populates PromptQueue.history on ComfyUI startup with the latest generated images."""
|
||||
"""Populate PromptQueue.history with the newest generated images."""
|
||||
if self._restored_on_launch and not force:
|
||||
return 0
|
||||
|
||||
started_at = time.time()
|
||||
if limit is None:
|
||||
limit = get_assets_restore_count()
|
||||
|
||||
@@ -121,119 +245,136 @@ class AssetsRestoreManager:
|
||||
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 []
|
||||
|
||||
images = self.scan_output_images(base_output_dir, limit=limit) 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
|
||||
|
||||
mtime_ms = int(item["mtime"] * 1000)
|
||||
duration = get_generation_time(filename, subfolder, item.get("mtime")) if show_generation_time else 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
|
||||
prompt_queue = getattr(server, "prompt_queue", None)
|
||||
|
||||
if prompt_queue is not None:
|
||||
with prompt_queue.mutex:
|
||||
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)
|
||||
output_node_id = None
|
||||
if prompt_dict:
|
||||
for nid, nval in prompt_dict.items():
|
||||
if isinstance(nval, dict) and "SaveImage" in 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 workflow_dict:
|
||||
extra_pnginfo["workflow"] = workflow_dict
|
||||
if positive_prompt:
|
||||
extra_pnginfo["positive_prompt"] = positive_prompt
|
||||
|
||||
mtime_ms = int(item["mtime"] * 1000)
|
||||
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": [
|
||||
["execution_start", {"timestamp": mtime_ms}],
|
||||
["execution_success", {"timestamp": mtime_ms}]
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
prompt_queue.history[prompt_id] = history_entry
|
||||
for entry in prepared:
|
||||
prompt_queue.history[entry["prompt_id"]] = entry["history_entry"]
|
||||
restored_count += 1
|
||||
|
||||
# Ingest into modern ComfyUI asset database if --enable-assets is active
|
||||
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 register_file_in_place:
|
||||
res = register_file_in_place(abs_path=filepath, name=filename, tags=["output"])
|
||||
if res and hasattr(res, "ref") and res.ref and hasattr(res.ref, "id"):
|
||||
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:
|
||||
set_reference_metadata(session, reference_id=res.ref.id, metadata={
|
||||
"jobId": prompt_id,
|
||||
"nodeId": output_node_id,
|
||||
"filename": filename,
|
||||
"subfolder": subfolder
|
||||
})
|
||||
session.commit()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
registered_count = self._ingest_assets(prepared)
|
||||
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,
|
||||
"total_scanned_images": len(images),
|
||||
"restored_count": restored_count
|
||||
"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),
|
||||
}
|
||||
self.restored_count = restored_count
|
||||
|
||||
if restored_count > 0:
|
||||
if server:
|
||||
|
||||
@@ -0,0 +1,273 @@
|
||||
"""
|
||||
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 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:
|
||||
return int(entry.get("completed", 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
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 = loaded
|
||||
except Exception:
|
||||
pass
|
||||
_DATA = data
|
||||
_DATA_MTIME = mtime
|
||||
return _DATA
|
||||
|
||||
|
||||
def _save_data(data):
|
||||
global _DATA_MTIME
|
||||
try:
|
||||
os.makedirs(USER_DIR, exist_ok=True)
|
||||
temp_file = GENERATION_TIMES_FILE + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, separators=(",", ":"), ensure_ascii=False)
|
||||
f.flush()
|
||||
try:
|
||||
os.fsync(f.fileno())
|
||||
except OSError:
|
||||
pass
|
||||
os.replace(temp_file, GENERATION_TIMES_FILE)
|
||||
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}")
|
||||
|
||||
|
||||
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):
|
||||
return
|
||||
if 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:
|
||||
return float(entry.get("duration"))
|
||||
except (TypeError, ValueError):
|
||||
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
|
||||
+132
-70
@@ -4,6 +4,9 @@ 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
|
||||
@@ -14,7 +17,8 @@ from .utils import (
|
||||
is_local_request,
|
||||
is_safe_path,
|
||||
sanitize_image_loader_folder,
|
||||
get_saturnnodes_user_dir
|
||||
get_saturnnodes_user_dir,
|
||||
get_env_setting
|
||||
)
|
||||
|
||||
IMAGE_CATEGORY = "🪐 SaturnNodes/Loaders"
|
||||
@@ -169,35 +173,31 @@ def save_prompts_cache(cache):
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Failed to save prompts cache: {e}")
|
||||
|
||||
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.")
|
||||
_PROMPTS_CACHE_LOCK = threading.RLock()
|
||||
IMAGE_THUMBNAILS_DIR = os.path.join(USER_DIR, "cache", "image_thumbs")
|
||||
|
||||
@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)
|
||||
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
|
||||
|
||||
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": {}})
|
||||
|
||||
image_extensions = {".png", ".jpg", ".jpeg", ".webp", ".PNG", ".JPG", ".JPEG", ".WEBP"}
|
||||
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 in image_extensions:
|
||||
if ext.lower() in image_extensions:
|
||||
full_path = os.path.normpath(os.path.join(root, file))
|
||||
if not is_safe_path(full_path):
|
||||
# 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)
|
||||
@@ -212,44 +212,102 @@ async def get_images_endpoint(request):
|
||||
names = []
|
||||
mapping = {}
|
||||
prompts = {}
|
||||
|
||||
cache = load_prompts_cache()
|
||||
cache_modified = False
|
||||
|
||||
# Extract prompt metadata for the most recent 128 images to keep cache small and fast
|
||||
METADATA_LIMIT = 128
|
||||
|
||||
for i, (rel_path, full_path, mtime) in enumerate(raw_files):
|
||||
names.append(rel_path)
|
||||
mapping[rel_path] = rel_path
|
||||
|
||||
if i < METADATA_LIMIT:
|
||||
|
||||
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:
|
||||
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] = {
|
||||
"mtime": mtime,
|
||||
"prompt": prompt_text
|
||||
}
|
||||
cache_modified = True
|
||||
except Exception:
|
||||
prompt_text = ""
|
||||
prompts[rel_path] = prompt_text
|
||||
else:
|
||||
prompts[rel_path] = ""
|
||||
|
||||
# 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:
|
||||
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] = {
|
||||
"mtime": mtime,
|
||||
"prompt": prompt_text
|
||||
}
|
||||
cache_modified = True
|
||||
except Exception as e:
|
||||
prompt_text = ""
|
||||
prompts[rel_path] = prompt_text
|
||||
else:
|
||||
prompts[rel_path] = ""
|
||||
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
|
||||
|
||||
# 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 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.")
|
||||
|
||||
if cache_modified:
|
||||
save_prompts_cache(cache)
|
||||
@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,
|
||||
@@ -266,25 +324,29 @@ async def get_thumbnail_endpoint(request):
|
||||
if not folder:
|
||||
folder = request.query.get("folder_path", "")
|
||||
resolved_folder = sanitize_image_loader_folder(folder)
|
||||
|
||||
|
||||
image_rel = request.query.get("image", "")
|
||||
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) or not os.path.exists(full_path) or not os.path.isfile(full_path):
|
||||
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)
|
||||
|
||||
@@ -297,13 +359,13 @@ async def get_full_image_endpoint(request):
|
||||
if not folder:
|
||||
folder = request.query.get("folder_path", "")
|
||||
resolved_folder = sanitize_image_loader_folder(folder)
|
||||
|
||||
|
||||
image_rel = request.query.get("image", "")
|
||||
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) or not os.path.exists(full_path) or not os.path.isfile(full_path):
|
||||
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"}
|
||||
|
||||
+43
-18
@@ -81,31 +81,36 @@ class TextLoraFinder:
|
||||
|
||||
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()
|
||||
for display_name, sys_path in mapping.items():
|
||||
if display_name in ["[ NONE ]", "[ RANDOM ]"] or not sys_path or sys_path == "[ NONE ]":
|
||||
continue
|
||||
matched_text = m.group(0).strip().lower()
|
||||
for display_name, sys_path in lookup.get(matched_text, []):
|
||||
if sys_path in seen_paths:
|
||||
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]
|
||||
|
||||
if (matched_text.lower() in [pretty_val.lower(), fname.lower(), fname_noext.lower(), display_name.lower()]):
|
||||
matches.append((m.start(), display_name, sys_path))
|
||||
seen_paths.add(sys_path)
|
||||
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 sys_path in seen_paths:
|
||||
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()
|
||||
@@ -116,11 +121,31 @@ class TextLoraFinder:
|
||||
|
||||
if not target_name:
|
||||
continue
|
||||
lookup.setdefault(target_name.lower(), []).append((display_name, sys_path))
|
||||
|
||||
pattern = r'(?<!\w)' + re.escape(target_name) + r'(?!\w)'
|
||||
for m in re.finditer(pattern, text_str, flags=re.IGNORECASE):
|
||||
matches.append((m.start(), display_name, sys_path))
|
||||
seen_paths.add(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, "")
|
||||
|
||||
+294
-130
@@ -5,6 +5,7 @@ import json
|
||||
import time
|
||||
import hashlib
|
||||
import asyncio
|
||||
import atexit
|
||||
import folder_paths
|
||||
import urllib.request
|
||||
import urllib.parse
|
||||
@@ -22,7 +23,8 @@ from .utils import (
|
||||
get_saturnnodes_user_dir,
|
||||
is_local_request,
|
||||
is_authenticated_local_request,
|
||||
is_safe_path
|
||||
is_safe_path,
|
||||
get_env_setting
|
||||
)
|
||||
|
||||
LORA_CATEGORY = "🪐 SaturnNodes/Loaders"
|
||||
@@ -39,34 +41,49 @@ CURRENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
ENV_FILE = os.path.join(USER_DIR, ".env")
|
||||
LORA_STATE_FILE = os.path.join(USER_DIR, "lora_loader_state.json")
|
||||
HASH_CACHE_FILE = os.path.join(USER_DIR, "lora_hash_cache.json")
|
||||
|
||||
_LORA_STATE_LOCK = threading.RLock()
|
||||
|
||||
def ensure_user_dir():
|
||||
os.makedirs(USER_DIR, exist_ok=True)
|
||||
|
||||
def _entry_timestamp(entry):
|
||||
if isinstance(entry, dict):
|
||||
try:
|
||||
return int(entry.get("last_updated", 0))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
return 0
|
||||
|
||||
def load_lora_cycle_state():
|
||||
ensure_user_dir()
|
||||
if os.path.exists(LORA_STATE_FILE):
|
||||
try:
|
||||
with open(LORA_STATE_FILE, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except Exception:
|
||||
pass
|
||||
return {}
|
||||
with _LORA_STATE_LOCK:
|
||||
if os.path.exists(LORA_STATE_FILE):
|
||||
try:
|
||||
with open(LORA_STATE_FILE, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
except Exception:
|
||||
pass
|
||||
return {}
|
||||
|
||||
def save_lora_cycle_state(state):
|
||||
ensure_user_dir()
|
||||
try:
|
||||
# Auto-prune abandoned state keys if count exceeds 100 entries
|
||||
if len(state) > 100:
|
||||
excess_keys = list(state.keys())[:-100]
|
||||
for k in excess_keys:
|
||||
del state[k]
|
||||
temp_file = LORA_STATE_FILE + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(state, f, indent=4, ensure_ascii=False)
|
||||
os.replace(temp_file, LORA_STATE_FILE)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error saving lora loader cycle state: {e}")
|
||||
with _LORA_STATE_LOCK:
|
||||
try:
|
||||
# Global prune: keep the 100 most recently used node/selection 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]
|
||||
temp_file = LORA_STATE_FILE + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(state, f, indent=4, ensure_ascii=False)
|
||||
os.replace(temp_file, LORA_STATE_FILE)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error saving lora loader cycle state: {e}")
|
||||
|
||||
def get_api_keys():
|
||||
civitai_key = os.getenv("CIVITAI_API_KEY", "")
|
||||
@@ -88,9 +105,18 @@ def get_api_keys():
|
||||
scraping_thread = None
|
||||
scraping_lock = threading.Lock()
|
||||
|
||||
def load_usage_data():
|
||||
usage_file = os.path.join(USER_DIR, "lora_usage.json")
|
||||
_USAGE_LOCK = threading.RLock()
|
||||
_USAGE_CACHE = None
|
||||
_USAGE_DIRTY = False
|
||||
_USAGE_TIMER = None
|
||||
_USAGE_FLUSH_DELAY = 2.0
|
||||
|
||||
def _usage_file_path():
|
||||
return os.path.join(USER_DIR, "lora_usage.json")
|
||||
|
||||
def _read_usage_data_from_disk():
|
||||
usage_data = {}
|
||||
usage_file = _usage_file_path()
|
||||
if os.path.exists(usage_file):
|
||||
try:
|
||||
with open(usage_file, "r", encoding="utf-8") as f:
|
||||
@@ -105,23 +131,69 @@ def load_usage_data():
|
||||
print(f"[SaturnNodes] Error loading usage data: {e}")
|
||||
return usage_data
|
||||
|
||||
def load_usage_data():
|
||||
global _USAGE_CACHE
|
||||
with _USAGE_LOCK:
|
||||
if _USAGE_CACHE is None:
|
||||
_USAGE_CACHE = _read_usage_data_from_disk()
|
||||
return dict(_USAGE_CACHE)
|
||||
|
||||
def flush_lora_usage():
|
||||
"""Write pending usage increments to disk (debounced and atomic)."""
|
||||
global _USAGE_DIRTY
|
||||
with _USAGE_LOCK:
|
||||
if not _USAGE_DIRTY or _USAGE_CACHE is None:
|
||||
return
|
||||
try:
|
||||
ensure_user_dir()
|
||||
usage_file = _usage_file_path()
|
||||
temp_file = usage_file + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(_USAGE_CACHE, f, indent=4, ensure_ascii=False)
|
||||
os.replace(temp_file, usage_file)
|
||||
_USAGE_DIRTY = False
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error updating lora usage: {e}")
|
||||
|
||||
def _schedule_usage_flush():
|
||||
global _USAGE_TIMER
|
||||
if _USAGE_TIMER is None or not _USAGE_TIMER.is_alive():
|
||||
_USAGE_TIMER = threading.Timer(_USAGE_FLUSH_DELAY, flush_lora_usage)
|
||||
_USAGE_TIMER.daemon = True
|
||||
_USAGE_TIMER.start()
|
||||
|
||||
atexit.register(flush_lora_usage)
|
||||
|
||||
def increment_lora_usage(lora_path):
|
||||
if not lora_path or lora_path == "[ NONE ]":
|
||||
return
|
||||
if get_env_setting("ENABLE_LORA_USAGE", "true").lower() != "true":
|
||||
if get_env_setting("ENABLE_LORA_USAGE", "true").lower() not in ["true", "1", "yes"]:
|
||||
return
|
||||
global _USAGE_CACHE, _USAGE_DIRTY
|
||||
try:
|
||||
usage_file = os.path.join(USER_DIR, "lora_usage.json")
|
||||
usage_data = load_usage_data()
|
||||
pretty_key = parse_pretty_name(lora_path)
|
||||
usage_data[pretty_key] = usage_data.get(pretty_key, 0) + 1
|
||||
temp_file = usage_file + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(usage_data, f, indent=4, ensure_ascii=False)
|
||||
os.replace(temp_file, usage_file)
|
||||
with _USAGE_LOCK:
|
||||
if _USAGE_CACHE is None:
|
||||
_USAGE_CACHE = _read_usage_data_from_disk()
|
||||
pretty_key = parse_pretty_name(lora_path)
|
||||
_USAGE_CACHE[pretty_key] = _USAGE_CACHE.get(pretty_key, 0) + 1
|
||||
_USAGE_DIRTY = True
|
||||
_schedule_usage_flush()
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error updating lora usage: {e}")
|
||||
|
||||
def reset_lora_usage_file():
|
||||
"""Clear usage counters, discarding any pending debounced writes."""
|
||||
global _USAGE_CACHE, _USAGE_DIRTY
|
||||
with _USAGE_LOCK:
|
||||
ensure_user_dir()
|
||||
usage_file = _usage_file_path()
|
||||
temp_file = usage_file + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump({}, f, indent=4)
|
||||
os.replace(temp_file, usage_file)
|
||||
_USAGE_CACHE = {}
|
||||
_USAGE_DIRTY = False
|
||||
|
||||
def find_preview_image(lora_relative_path):
|
||||
lora_full_path = folder_paths.get_full_path("loras", lora_relative_path)
|
||||
if not lora_full_path:
|
||||
@@ -141,6 +213,73 @@ def get_file_sha256(filepath):
|
||||
h.update(chunk)
|
||||
return h.hexdigest().upper()
|
||||
|
||||
_HASH_CACHE_LOCK = threading.RLock()
|
||||
_HASH_CACHE = None
|
||||
|
||||
def _load_hash_cache():
|
||||
global _HASH_CACHE
|
||||
if _HASH_CACHE is not None:
|
||||
return _HASH_CACHE
|
||||
_HASH_CACHE = {}
|
||||
try:
|
||||
if os.path.exists(HASH_CACHE_FILE):
|
||||
with open(HASH_CACHE_FILE, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, dict):
|
||||
_HASH_CACHE = data
|
||||
except Exception:
|
||||
_HASH_CACHE = {}
|
||||
return _HASH_CACHE
|
||||
|
||||
def _save_hash_cache():
|
||||
try:
|
||||
ensure_user_dir()
|
||||
temp_file = HASH_CACHE_FILE + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(_HASH_CACHE, f, indent=2, ensure_ascii=False)
|
||||
os.replace(temp_file, HASH_CACHE_FILE)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def get_file_sha256_cached(filepath):
|
||||
"""Hash a LoRA once and reuse it while mtime/size stay unchanged."""
|
||||
try:
|
||||
before = os.stat(filepath)
|
||||
except OSError:
|
||||
return get_file_sha256(filepath)
|
||||
|
||||
cache_key = os.path.normcase(os.path.abspath(filepath))
|
||||
with _HASH_CACHE_LOCK:
|
||||
cache = _load_hash_cache()
|
||||
entry = cache.get(cache_key)
|
||||
if (isinstance(entry, dict)
|
||||
and entry.get("mtime_ns") == before.st_mtime_ns
|
||||
and entry.get("size") == before.st_size
|
||||
and entry.get("sha256")):
|
||||
return entry["sha256"]
|
||||
|
||||
digest = get_file_sha256(filepath)
|
||||
|
||||
try:
|
||||
after = os.stat(filepath)
|
||||
if after.st_mtime_ns == before.st_mtime_ns and after.st_size == before.st_size:
|
||||
with _HASH_CACHE_LOCK:
|
||||
cache = _load_hash_cache()
|
||||
cache[cache_key] = {
|
||||
"mtime_ns": before.st_mtime_ns,
|
||||
"size": before.st_size,
|
||||
"sha256": digest,
|
||||
"last_updated": int(time.time()),
|
||||
}
|
||||
if len(cache) > 1000:
|
||||
keep = set(sorted(cache.keys(), key=lambda k: int(cache[k].get("last_updated", 0)), reverse=True)[:1000])
|
||||
for k in [k for k in cache.keys() if k not in keep]:
|
||||
del cache[k]
|
||||
_save_hash_cache()
|
||||
except OSError:
|
||||
pass
|
||||
return digest
|
||||
|
||||
def is_safe_external_image_url(url_str):
|
||||
"""
|
||||
Validates that a scraped image URL uses a safe web scheme (https/http)
|
||||
@@ -167,17 +306,6 @@ def is_safe_external_image_url(url_str):
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def get_env_setting(key, default_val):
|
||||
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(f"{key}="):
|
||||
return line.strip().split("=", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
return default_val
|
||||
|
||||
def scrape_missing_images_sync():
|
||||
try:
|
||||
civitai_enabled = get_env_setting("ENABLE_CIVITAI_SCRAPING", "false").lower() in ["true", "1", "yes"]
|
||||
@@ -218,7 +346,7 @@ def scrape_missing_images_sync():
|
||||
|
||||
if civitai_enabled:
|
||||
try:
|
||||
file_hash = get_file_sha256(lora_path)
|
||||
file_hash = get_file_sha256_cached(lora_path)
|
||||
civitai_url = f"https://civitai.com/api/v1/model-versions/by-hash/{file_hash}"
|
||||
headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)'}
|
||||
if civitai_key:
|
||||
@@ -377,7 +505,7 @@ async def get_loras_endpoint(request):
|
||||
if scrape_on_new:
|
||||
start_async_scraping()
|
||||
|
||||
usage_enabled = get_env_setting("ENABLE_LORA_USAGE", "true").lower() == "true"
|
||||
usage_enabled = get_env_setting("ENABLE_LORA_USAGE", "true").lower() in ["true", "1", "yes"]
|
||||
usage_by_path = {}
|
||||
mtime_by_path = {}
|
||||
|
||||
@@ -504,12 +632,7 @@ async def reset_lora_usage_endpoint(request):
|
||||
if not is_authenticated_local_request(request):
|
||||
return web.json_response({"error": "Forbidden: Local authenticated access only"}, status=403)
|
||||
try:
|
||||
ensure_user_dir()
|
||||
usage_file = os.path.join(USER_DIR, "lora_usage.json")
|
||||
tmp_file = usage_file + ".tmp"
|
||||
with open(tmp_file, "w", encoding="utf-8") as f:
|
||||
json.dump({}, f, indent=4)
|
||||
os.replace(tmp_file, usage_file)
|
||||
reset_lora_usage_file()
|
||||
return web.json_response({"status": "ok"})
|
||||
except Exception as e:
|
||||
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
||||
@@ -664,16 +787,13 @@ class FolderLoraLoaderPretty(FolderLoraLoader):
|
||||
out_name = format_lora_output_name(resolved_path, display_name, output_format=active_format, custom_regex=custom_regex)
|
||||
|
||||
if strength_model == 0 and strength_clip == 0:
|
||||
increment_lora_usage(resolved_path)
|
||||
return (model, clip, out_name)
|
||||
|
||||
if not resolved_path:
|
||||
return (model, clip, out_name)
|
||||
|
||||
lora_path = folder_paths.get_full_path("loras", resolved_path)
|
||||
if strength_model == 0 and strength_clip == 0:
|
||||
increment_lora_usage(resolved_path)
|
||||
return (model, clip, out_name)
|
||||
|
||||
if not lora_path or not os.path.exists(lora_path):
|
||||
return (model, clip, out_name)
|
||||
|
||||
@@ -688,14 +808,127 @@ def get_current_prompt_id():
|
||||
try:
|
||||
from server import PromptServer
|
||||
server = PromptServer.instance
|
||||
if hasattr(server, "prompt_queue") and server.prompt_queue.currently_running:
|
||||
for item in server.prompt_queue.currently_running.values():
|
||||
if hasattr(server, "prompt_queue"):
|
||||
queue = server.prompt_queue
|
||||
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_lora_assignment(assigned, prompt_id, item, limit=50):
|
||||
"""Insert/refresh a prompt→LoRA mapping, keeping the most recent entries."""
|
||||
if not prompt_id:
|
||||
return
|
||||
assigned.pop(prompt_id, None)
|
||||
assigned[prompt_id] = item
|
||||
while len(assigned) > limit:
|
||||
assigned.pop(next(iter(assigned)))
|
||||
|
||||
def select_a_lora(mode, unique_id, valid_choices, prompt_id=None):
|
||||
"""
|
||||
Choose one LoRA for Sequential / Random / Random (No Replace), persisting
|
||||
cycle position, pools and crash-recovery prompt assignments in one atomic
|
||||
state update. Returns None when there is nothing to select.
|
||||
"""
|
||||
valid_choices = [item for item in valid_choices if item and item != "[ NONE ]"]
|
||||
if not valid_choices:
|
||||
return None
|
||||
|
||||
choices_hash = hashlib.sha256(
|
||||
json.dumps(sorted(valid_choices), ensure_ascii=False).encode("utf-8")
|
||||
).hexdigest()[:16]
|
||||
prefix = {
|
||||
"Random": "rand",
|
||||
"Sequential": "seq",
|
||||
"Random (No Replace)": "noreplace",
|
||||
}.get(mode)
|
||||
if not prefix:
|
||||
return None
|
||||
|
||||
state_key = f"{prefix}_{unique_id}_{choices_hash}"
|
||||
|
||||
with _LORA_STATE_LOCK:
|
||||
state = load_lora_cycle_state()
|
||||
node_state = state.get(state_key)
|
||||
if isinstance(node_state, dict):
|
||||
raw_assigned = node_state.get("assigned", {})
|
||||
assigned = dict(raw_assigned) if isinstance(raw_assigned, dict) else {}
|
||||
raw_pool = node_state.get("pool", [])
|
||||
pool = list(raw_pool) if isinstance(raw_pool, (list, tuple)) else []
|
||||
last_selected = node_state.get("last")
|
||||
try:
|
||||
index = max(0, int(node_state.get("index", 0)))
|
||||
except (TypeError, ValueError):
|
||||
index = 0
|
||||
elif isinstance(node_state, list):
|
||||
# Legacy Random (No Replace) format stored a bare pool list.
|
||||
assigned = {}
|
||||
pool = list(node_state)
|
||||
last_selected = None
|
||||
index = 0
|
||||
else:
|
||||
assigned = {}
|
||||
pool = []
|
||||
last_selected = None
|
||||
index = 0
|
||||
|
||||
if prompt_id and prompt_id in assigned:
|
||||
# Retry/crash recovery: reproduce the exact same selection without
|
||||
# advancing the cycle or consuming another pool entry.
|
||||
selected = assigned[prompt_id]
|
||||
if selected not in valid_choices:
|
||||
selected = None
|
||||
else:
|
||||
selected = None
|
||||
|
||||
if selected is None:
|
||||
if mode == "Sequential":
|
||||
selected = valid_choices[index % len(valid_choices)]
|
||||
index += 1
|
||||
elif mode == "Random":
|
||||
import random
|
||||
selected = random.choice(valid_choices)
|
||||
elif mode == "Random (No Replace)":
|
||||
import random
|
||||
pool = [x for x in pool if x in valid_choices]
|
||||
if not pool:
|
||||
pool = list(valid_choices)
|
||||
# Avoid repeating the last item of the previous cycle
|
||||
# back-to-back while still visiting every item once.
|
||||
if len(pool) > 1 and last_selected in pool:
|
||||
random.shuffle(pool)
|
||||
if pool[0] == last_selected:
|
||||
swap_index = random.randrange(1, len(pool))
|
||||
pool[0], pool[swap_index] = pool[swap_index], pool[0]
|
||||
else:
|
||||
random.shuffle(pool)
|
||||
selected = pool.pop(0)
|
||||
last_selected = selected
|
||||
|
||||
if prompt_id and selected:
|
||||
_remember_lora_assignment(assigned, prompt_id, selected)
|
||||
|
||||
new_entry = {
|
||||
"assigned": assigned,
|
||||
"last_updated": int(time.time()),
|
||||
}
|
||||
if mode == "Sequential":
|
||||
new_entry["index"] = index
|
||||
if mode == "Random (No Replace)":
|
||||
new_entry["pool"] = pool
|
||||
new_entry["last"] = last_selected
|
||||
state[state_key] = new_entry
|
||||
save_lora_cycle_state(state)
|
||||
return selected
|
||||
|
||||
class VisualLoraLoader(FolderLoraLoaderPretty):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -751,82 +984,14 @@ class VisualLoraLoader(FolderLoraLoaderPretty):
|
||||
current_pid = get_current_prompt_id()
|
||||
if _selection_mode in ["Random", "Sequential", "Random (No Replace)"] and len(loras_to_load) > 0:
|
||||
valid_choices = [item for item in loras_to_load if item != "[ NONE ]" and item]
|
||||
if valid_choices:
|
||||
if _selection_mode == "Random":
|
||||
import random
|
||||
state = load_lora_cycle_state()
|
||||
text_hash = hashlib.sha256(json.dumps(sorted(valid_choices)).encode('utf-8')).hexdigest()[:16]
|
||||
state_key = f"rand_{unique_id}_{text_hash}"
|
||||
node_state = state.get(state_key, {})
|
||||
assigned = node_state.get("assigned", {}) if isinstance(node_state, dict) else {}
|
||||
if current_pid and current_pid in assigned:
|
||||
selected_item = assigned[current_pid]
|
||||
else:
|
||||
selected_item = random.choice(valid_choices)
|
||||
if current_pid:
|
||||
assigned[current_pid] = selected_item
|
||||
state[state_key] = {"assigned": assigned}
|
||||
save_lora_cycle_state(state)
|
||||
loras_to_load = [selected_item]
|
||||
elif _selection_mode == "Sequential":
|
||||
state = load_lora_cycle_state()
|
||||
text_hash = hashlib.sha256(json.dumps(sorted(valid_choices)).encode('utf-8')).hexdigest()[:16]
|
||||
state_key = f"seq_{unique_id}_{text_hash}"
|
||||
node_state = state.get(state_key, {})
|
||||
if isinstance(node_state, dict):
|
||||
idx = int(node_state.get("index", 0))
|
||||
assigned = node_state.get("assigned", {})
|
||||
else:
|
||||
idx = int(node_state) if isinstance(node_state, int) else 0
|
||||
assigned = {}
|
||||
|
||||
if current_pid and current_pid in assigned:
|
||||
selected_item = assigned[current_pid]
|
||||
else:
|
||||
selected_item = valid_choices[idx % len(valid_choices)]
|
||||
if current_pid:
|
||||
assigned[current_pid] = selected_item
|
||||
state[state_key] = {
|
||||
"index": idx + 1,
|
||||
"assigned": assigned
|
||||
}
|
||||
save_lora_cycle_state(state)
|
||||
loras_to_load = [selected_item]
|
||||
elif _selection_mode == "Random (No Replace)":
|
||||
state = load_lora_cycle_state()
|
||||
text_hash = hashlib.sha256(json.dumps(sorted(valid_choices)).encode('utf-8')).hexdigest()[:16]
|
||||
state_key = f"noreplace_{unique_id}_{text_hash}"
|
||||
node_state = state.get(state_key, {})
|
||||
if isinstance(node_state, dict):
|
||||
pool = node_state.get("pool", [])
|
||||
assigned = node_state.get("assigned", {})
|
||||
else:
|
||||
pool = node_state if isinstance(node_state, list) else []
|
||||
assigned = {}
|
||||
|
||||
if current_pid and current_pid in assigned:
|
||||
selected_item = assigned[current_pid]
|
||||
else:
|
||||
pool = [x for x in pool if x in valid_choices]
|
||||
if not pool:
|
||||
import random
|
||||
pool = list(valid_choices)
|
||||
random.shuffle(pool)
|
||||
selected_item = pool.pop(0)
|
||||
if current_pid:
|
||||
assigned[current_pid] = selected_item
|
||||
state[state_key] = {
|
||||
"pool": pool,
|
||||
"assigned": assigned
|
||||
}
|
||||
save_lora_cycle_state(state)
|
||||
loras_to_load = [selected_item]
|
||||
else:
|
||||
loras_to_load = []
|
||||
selected_item = select_a_lora(_selection_mode, unique_id, valid_choices, current_pid)
|
||||
loras_to_load = [selected_item] if selected_item else []
|
||||
|
||||
current_model = model
|
||||
current_clip = clip
|
||||
loaded_names = []
|
||||
# Resolve the folder mapping once per run instead of once per selected LoRA.
|
||||
lora_mapping = get_filtered_loras_mapping(folder, pretty=True)
|
||||
|
||||
for item in loras_to_load:
|
||||
if item == "[ NONE ]" or not item:
|
||||
@@ -836,8 +1001,7 @@ class VisualLoraLoader(FolderLoraLoaderPretty):
|
||||
if " - " in item:
|
||||
display_name = item.split(" - ", 1)[1]
|
||||
|
||||
mapping = get_filtered_loras_mapping(folder, pretty=True)
|
||||
resolved_path = mapping.get(item)
|
||||
resolved_path = lora_mapping.get(item)
|
||||
|
||||
if not resolved_path or resolved_path == "[ NONE ]":
|
||||
all_loras = folder_paths.get_filename_list("loras")
|
||||
|
||||
+208
-86
@@ -1,31 +1,25 @@
|
||||
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
|
||||
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")
|
||||
ENV_FILE = os.path.join(USER_DIR, ".env")
|
||||
|
||||
def get_env_setting(key, default_val):
|
||||
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(f"{key}="):
|
||||
return line.strip().split("=", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
return default_val
|
||||
_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()
|
||||
@@ -34,45 +28,84 @@ def is_clear_prompt_iterator_on_launch():
|
||||
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()
|
||||
if os.path.exists(STATE_FILE):
|
||||
with _STATE_LOCK:
|
||||
mtime = None
|
||||
try:
|
||||
with open(STATE_FILE, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Warning reading prompt iterator state: {e}")
|
||||
return {}
|
||||
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()
|
||||
try:
|
||||
# Auto-prune abandoned state keys if count exceeds 100 entries
|
||||
if len(state) > 100:
|
||||
excess_keys = list(state.keys())[:-100]
|
||||
for k in excess_keys:
|
||||
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)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error saving prompt iterator state: {e}")
|
||||
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()
|
||||
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)
|
||||
print("[SaturnNodes] Prompt iterator state cleared.")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes] Error clearing prompt iterator state: {e}")
|
||||
return False
|
||||
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():
|
||||
@@ -90,14 +123,7 @@ try:
|
||||
data = await request.json()
|
||||
node_id = str(data.get("node_id", ""))
|
||||
state = load_state()
|
||||
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
|
||||
else:
|
||||
state[k] = {"index": 0, "total": 0}
|
||||
matched = True
|
||||
matched = reset_node_state(state, node_id)
|
||||
save_state(state)
|
||||
if node_id and node_id != "all":
|
||||
progress_payload = {
|
||||
@@ -185,12 +211,31 @@ def parse_prompt_blocks(text_str, separator, custom_regex=""):
|
||||
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") and queue.currently_running:
|
||||
if hasattr(queue, "currently_running"):
|
||||
mutex = getattr(queue, "mutex", None)
|
||||
if mutex:
|
||||
with mutex:
|
||||
@@ -204,6 +249,28 @@ def get_current_prompt_id():
|
||||
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):
|
||||
@@ -225,6 +292,12 @@ class PromptQueueIterator:
|
||||
},
|
||||
"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": {
|
||||
@@ -232,8 +305,8 @@ class PromptQueueIterator:
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "INT")
|
||||
RETURN_NAMES = ("prompt", "remaining_text", "remaining_count")
|
||||
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."
|
||||
@@ -246,6 +319,7 @@ class PromptQueueIterator:
|
||||
def process_queue(
|
||||
self,
|
||||
pop_mode="Sequential (Loop on End)",
|
||||
track_by="Text & Node",
|
||||
separator=">1 Empty Line",
|
||||
text="",
|
||||
custom_regex="",
|
||||
@@ -257,17 +331,32 @@ class PromptQueueIterator:
|
||||
raw_input = text if text else (prompt if prompt is not None else prompt_text)
|
||||
text_str = str(raw_input) if raw_input is not None else ""
|
||||
if not text_str.strip():
|
||||
return ("", "", 0)
|
||||
return ("", "", 0, False)
|
||||
|
||||
original_blocks = parse_prompt_blocks(text_str, separator, custom_regex)
|
||||
total_count = len(original_blocks)
|
||||
if total_count == 0:
|
||||
return ("", "", 0)
|
||||
return ("", "", 0, False)
|
||||
|
||||
# Calculate stable hash of the complete input text
|
||||
# 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]
|
||||
state_key = f"node_{unique_id}_{text_hash}"
|
||||
|
||||
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)
|
||||
@@ -276,67 +365,100 @@ class PromptQueueIterator:
|
||||
|
||||
current_index = 0
|
||||
assigned_prompts = {}
|
||||
shuffled_order = []
|
||||
if isinstance(node_state, dict):
|
||||
current_index = int(node_state.get("index", 0))
|
||||
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 = node_state
|
||||
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)
|
||||
if current_prompt_id and current_prompt_id in assigned_prompts:
|
||||
idx = assigned_prompts[current_prompt_id] % total_count
|
||||
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]
|
||||
display_run = idx + 1
|
||||
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
|
||||
display_run = idx + 1
|
||||
elif pop_mode in ["Sequential (Stop on End)"]:
|
||||
elif is_stop_mode:
|
||||
if current_index >= total_count:
|
||||
idx = total_count - 1
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = total_count
|
||||
display_run = total_count
|
||||
else:
|
||||
idx = current_index
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = current_index + 1
|
||||
display_run = idx + 1
|
||||
elif pop_mode in ["Random (Keep)", "Random (Delete)"]:
|
||||
elif is_random_keep:
|
||||
idx = random.randint(0, total_count - 1)
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = current_index + 1
|
||||
display_run = (current_index % total_count) + 1
|
||||
elif pop_mode in ["Random (Cycle)"]:
|
||||
shuffled_order = list(node_state.get("shuffled_order", [])) if isinstance(node_state, dict) else []
|
||||
elif is_random_cycle:
|
||||
cycle_index = current_index % total_count
|
||||
if len(shuffled_order) != total_count or (cycle_index == 0 and current_index > 0) or not shuffled_order:
|
||||
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
|
||||
display_run = cycle_index + 1
|
||||
else: # Default fallback
|
||||
idx = current_index % total_count
|
||||
selected_prompt = original_blocks[idx]
|
||||
next_index = (current_index + 1) % total_count
|
||||
display_run = idx + 1
|
||||
|
||||
if current_prompt_id:
|
||||
assigned_prompts[current_prompt_id] = idx
|
||||
if len(assigned_prompts) > 50:
|
||||
for old_k in list(assigned_prompts.keys())[:-50]:
|
||||
del assigned_prompts[old_k]
|
||||
_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
|
||||
|
||||
# Calculate remaining items
|
||||
remaining_count = max(0, total_count - display_run)
|
||||
sep_lower = str(separator).lower() if separator else ""
|
||||
if "newline" in sep_lower:
|
||||
join_delim = "\n"
|
||||
@@ -349,7 +471,7 @@ class PromptQueueIterator:
|
||||
join_delim = "\n\n"
|
||||
else:
|
||||
join_delim = "\n\n"
|
||||
remaining_text = join_delim.join(original_blocks[display_run:]) if display_run < total_count else ""
|
||||
remaining_text = join_delim.join(remaining_blocks)
|
||||
|
||||
# Update state dictionary
|
||||
new_entry = {
|
||||
@@ -360,15 +482,15 @@ class PromptQueueIterator:
|
||||
"preview": selected_prompt[:60].replace('\n', ' '),
|
||||
"assigned_prompts": assigned_prompts
|
||||
}
|
||||
if "shuffled_order" in locals():
|
||||
if is_random_cycle:
|
||||
new_entry["shuffled_order"] = shuffled_order
|
||||
state[state_key] = new_entry
|
||||
|
||||
# Keep at most 3 most recent text hashes per node_id
|
||||
# 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: state[k].get("last_updated", 0) if isinstance(state[k], dict) else 0)
|
||||
node_keys.sort(key=lambda k: _entry_timestamp(state[k]))
|
||||
for old_k in node_keys[:-3]:
|
||||
del state[old_k]
|
||||
|
||||
@@ -386,4 +508,4 @@ class PromptQueueIterator:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return (selected_prompt, remaining_text, remaining_count)
|
||||
return (selected_prompt, remaining_text, remaining_count, exhausted)
|
||||
|
||||
+255
-46
@@ -4,30 +4,25 @@ import secrets
|
||||
import time
|
||||
import threading
|
||||
import asyncio
|
||||
import atexit
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
import nodes
|
||||
from .tray_icon import TrayIconManager, is_tray_icon_enabled
|
||||
from .assets_restore import assets_restore_manager, is_assets_restore_enabled, get_assets_restore_count
|
||||
from .utils import get_saturnnodes_user_dir, is_local_request, is_authenticated_local_request
|
||||
from .generation_times import patch_prompt_queue as patch_generation_times
|
||||
from .utils import get_saturnnodes_user_dir, is_local_request, is_authenticated_local_request, get_env_setting
|
||||
|
||||
QUEUE_CATEGORY = "🪐 SaturnNodes/Queue"
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
USER_DIR = get_saturnnodes_user_dir()
|
||||
PERSISTENT_FILE = os.path.join(USER_DIR, "persistent_queue.json")
|
||||
ENV_FILE = os.path.join(USER_DIR, ".env")
|
||||
|
||||
def get_env_setting(key, default_val):
|
||||
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(f"{key}="):
|
||||
return line.strip().split("=", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
return default_val
|
||||
# Versioned envelope for persistent_queue.json. Bump this when the persisted
|
||||
# entry shape changes and add a migration branch in
|
||||
# PersistentQueueManager._normalize_persisted_payload().
|
||||
PERSISTENT_QUEUE_SCHEMA_VERSION = 1
|
||||
|
||||
def is_persistent_queue_enabled():
|
||||
val = get_env_setting("ENABLE_PERSISTENT_QUEUE", "true").lower()
|
||||
@@ -41,9 +36,25 @@ class PauseQueueManager:
|
||||
def __init__(self):
|
||||
default_state = get_env_setting("DEFAULT_PAUSE_STATE", "Running")
|
||||
default_mode = get_env_setting("DEFAULT_PAUSE_MODE", "after_finish")
|
||||
|
||||
self.paused = (default_state.lower() != "running")
|
||||
self.mode = "instantly" if "instant" in default_mode.lower() else "after_finish"
|
||||
self.persist_state = get_env_setting("PERSIST_PAUSE_STATE", "false").lower() in ["true", "1", "yes"]
|
||||
self.state_file = os.path.join(USER_DIR, "pause_state.json")
|
||||
|
||||
paused = (default_state.lower() != "running")
|
||||
mode = "instantly" if "instant" in default_mode.lower() else "after_finish"
|
||||
|
||||
if self.persist_state and os.path.exists(self.state_file):
|
||||
try:
|
||||
with open(self.state_file, "r", encoding="utf-8") as f:
|
||||
saved = json.load(f)
|
||||
if isinstance(saved, dict):
|
||||
paused = bool(saved.get("paused", paused))
|
||||
saved_mode = str(saved.get("mode", mode))
|
||||
mode = "instantly" if "instant" in saved_mode.lower() else "after_finish"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.paused = paused
|
||||
self.mode = mode
|
||||
self.is_waiting = self.paused
|
||||
self.event = threading.Event()
|
||||
if not self.paused:
|
||||
@@ -52,6 +63,18 @@ class PauseQueueManager:
|
||||
self.event.clear()
|
||||
self._patched = False
|
||||
|
||||
def _persist_pause_state(self):
|
||||
if not self.persist_state:
|
||||
return
|
||||
try:
|
||||
os.makedirs(USER_DIR, exist_ok=True)
|
||||
temp_file = self.state_file + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump({"paused": bool(self.paused), "mode": self.mode}, f)
|
||||
os.replace(temp_file, self.state_file)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def patch_all(self):
|
||||
if self._patched:
|
||||
return
|
||||
@@ -68,6 +91,8 @@ class PauseQueueManager:
|
||||
original_get = queue.get
|
||||
|
||||
def patched_get(*args, **kwargs):
|
||||
caller_timeout = args[0] if args else kwargs.get("timeout", None)
|
||||
caller_blocks = caller_timeout is None
|
||||
while True:
|
||||
# 1. If queue is paused, wait BEFORE popping any item from the queue!
|
||||
# This guarantees that ComfyUI never sets currently_running, never sends
|
||||
@@ -99,6 +124,14 @@ class PauseQueueManager:
|
||||
if item is not None:
|
||||
return item
|
||||
|
||||
# Honor a finite caller timeout. Returning None lets ComfyUI's main
|
||||
# loop run its scheduled GC pass and avoids a busy spin when a
|
||||
# timeout of 0.0 is requested.
|
||||
if self.paused:
|
||||
continue
|
||||
if not caller_blocks:
|
||||
return None
|
||||
|
||||
queue.get = patched_get
|
||||
|
||||
if hasattr(server, "send_sync") and callable(server.send_sync):
|
||||
@@ -120,6 +153,7 @@ class PauseQueueManager:
|
||||
self.paused = False
|
||||
self.is_waiting = False
|
||||
self.event.set()
|
||||
self._persist_pause_state()
|
||||
self.notify_clients()
|
||||
if hasattr(mm, "throw_exception_if_processing_interrupted"):
|
||||
mm.throw_exception_if_processing_interrupted()
|
||||
@@ -178,6 +212,7 @@ class PauseQueueManager:
|
||||
print("[SaturnNodes PauseQueue] Continue pressed: Resuming execution.")
|
||||
|
||||
self.notify_clients()
|
||||
self._persist_pause_state()
|
||||
|
||||
try:
|
||||
server = getattr(PromptServer, "instance", None)
|
||||
@@ -205,6 +240,7 @@ class PauseQueueManager:
|
||||
def set_mode(self, mode):
|
||||
self.mode = mode
|
||||
print(f"[SaturnNodes PauseQueue] Selected pause mode: '{self.mode}'")
|
||||
self._persist_pause_state()
|
||||
self.notify_clients()
|
||||
|
||||
def notify_clients(self):
|
||||
@@ -224,61 +260,198 @@ class PauseQueueManager:
|
||||
pass
|
||||
|
||||
class PersistentQueueManager:
|
||||
MAX_BATCH_META = 500
|
||||
MAX_RECENTLY_FINISHED = 500
|
||||
SAVE_DEBOUNCE_SECONDS = 0.25
|
||||
|
||||
def __init__(self):
|
||||
self.lock = threading.RLock()
|
||||
self.persistent_items = []
|
||||
self.batch_meta = {}
|
||||
self.batch_meta_updated = {}
|
||||
self.recently_finished = {}
|
||||
self.is_restoring = False
|
||||
self._patched = False
|
||||
self.has_claimed_once = False
|
||||
self._save_timer = None
|
||||
self._save_dirty = False
|
||||
# Set when persistent_queue.json was written by a newer SaturnNodes.
|
||||
# In that case the file is read-only for this process so a downgrade
|
||||
# cannot destroy a queue whose entry shape we do not understand.
|
||||
self._read_only_newer_schema = False
|
||||
self.load_from_file()
|
||||
|
||||
@staticmethod
|
||||
def _normalize_persisted_payload(data):
|
||||
"""Return ``(items, needs_rewrite, is_foreign)`` for persisted payloads.
|
||||
|
||||
- v1 files (pre-2.6.2) are a bare list and are migrated on load.
|
||||
- Current files are ``{"schema_version": 1, "items": [...]}``.
|
||||
- A file written by a newer SaturnNodes is reported as foreign so the
|
||||
caller can keep it read-only instead of destroying a queue it cannot
|
||||
understand.
|
||||
"""
|
||||
def is_valid_item(x):
|
||||
return (
|
||||
isinstance(x, dict)
|
||||
and "prompt_id" in x
|
||||
and isinstance(x.get("item"), (list, tuple))
|
||||
)
|
||||
|
||||
if isinstance(data, dict):
|
||||
version = data.get("schema_version")
|
||||
items = data.get("items")
|
||||
if not isinstance(items, list):
|
||||
return [], False, False
|
||||
if isinstance(version, int) and version > PERSISTENT_QUEUE_SCHEMA_VERSION:
|
||||
print(
|
||||
"[SaturnNodes PersistentQueue] persistent_queue.json was written by a newer "
|
||||
f"SaturnNodes (schema {version} > {PERSISTENT_QUEUE_SCHEMA_VERSION}); "
|
||||
"leaving it untouched."
|
||||
)
|
||||
return [], False, True
|
||||
return (
|
||||
[x for x in items if is_valid_item(x)],
|
||||
version != PERSISTENT_QUEUE_SCHEMA_VERSION,
|
||||
False,
|
||||
)
|
||||
|
||||
if isinstance(data, list):
|
||||
# Legacy bare-list format.
|
||||
return [x for x in data if is_valid_item(x)], True, False
|
||||
|
||||
return [], False, False
|
||||
|
||||
def load_from_file(self):
|
||||
with self.lock:
|
||||
if os.path.exists(PERSISTENT_FILE):
|
||||
try:
|
||||
with open(PERSISTENT_FILE, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, list):
|
||||
self.persistent_items = [
|
||||
x for x in data
|
||||
if isinstance(x, dict) and "prompt_id" in x and isinstance(x.get("item"), (list, tuple))
|
||||
]
|
||||
for x in self.persistent_items:
|
||||
if "batch_info" in x and "prompt_id" in x:
|
||||
self.batch_meta[str(x["prompt_id"])] = x["batch_info"]
|
||||
else:
|
||||
self.persistent_items = []
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes PersistentQueue] Error loading persistent queue file: {e}")
|
||||
self.persistent_items = []
|
||||
else:
|
||||
if not os.path.exists(PERSISTENT_FILE):
|
||||
self.persistent_items = []
|
||||
self._prune_batch_meta_locked()
|
||||
return
|
||||
try:
|
||||
with open(PERSISTENT_FILE, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
items, needs_rewrite, is_foreign = self._normalize_persisted_payload(data)
|
||||
self._read_only_newer_schema = is_foreign
|
||||
self.persistent_items = items
|
||||
for x in self.persistent_items:
|
||||
if "batch_info" in x and "prompt_id" in x:
|
||||
pid = str(x["prompt_id"])
|
||||
self.batch_meta[pid] = x["batch_info"]
|
||||
self.batch_meta_updated[pid] = time.time()
|
||||
if needs_rewrite:
|
||||
# Rewrite once so the legacy bare-list shape is not read
|
||||
# forever. _write_locked() is a no-op while persistence is
|
||||
# disabled, in which case the migration is retried later.
|
||||
self._save_dirty = True
|
||||
self._write_locked()
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes PersistentQueue] Error loading persistent queue file: {e}")
|
||||
self.persistent_items = []
|
||||
self._prune_batch_meta_locked()
|
||||
|
||||
def _prune_batch_meta_locked(self):
|
||||
if len(self.batch_meta) <= self.MAX_BATCH_META:
|
||||
return
|
||||
active_ids = {str(x.get("prompt_id")) for x in self.persistent_items}
|
||||
removable = sorted(
|
||||
(pid for pid in self.batch_meta if pid not in active_ids),
|
||||
key=lambda pid: self.batch_meta_updated.get(pid, 0),
|
||||
)
|
||||
for pid in removable:
|
||||
if len(self.batch_meta) <= self.MAX_BATCH_META:
|
||||
break
|
||||
self.batch_meta.pop(pid, None)
|
||||
self.batch_meta_updated.pop(pid, None)
|
||||
|
||||
def _mark_finished_locked(self, pid_str):
|
||||
self.recently_finished.pop(pid_str, None)
|
||||
self.recently_finished[pid_str] = time.time()
|
||||
while len(self.recently_finished) > self.MAX_RECENTLY_FINISHED:
|
||||
self.recently_finished.pop(next(iter(self.recently_finished)))
|
||||
|
||||
def set_batch_info(self, prompt_id, batch_info):
|
||||
if not prompt_id:
|
||||
return
|
||||
pid_str = str(prompt_id)
|
||||
with self.lock:
|
||||
# Late syncs from another client must not resurrect metadata for a
|
||||
# prompt that already finished (the id is single-use).
|
||||
if pid_str in self.recently_finished:
|
||||
return
|
||||
self.batch_meta[pid_str] = batch_info
|
||||
self.batch_meta_updated[pid_str] = time.time()
|
||||
for item in self.persistent_items:
|
||||
if str(item.get("prompt_id")) == pid_str:
|
||||
item["batch_info"] = batch_info
|
||||
break
|
||||
self._prune_batch_meta_locked()
|
||||
if is_persistent_queue_enabled():
|
||||
self.save_to_file()
|
||||
|
||||
def save_to_file(self):
|
||||
def save_to_file(self, immediate=False):
|
||||
if not is_persistent_queue_enabled():
|
||||
return
|
||||
with self.lock:
|
||||
try:
|
||||
temp_file = PERSISTENT_FILE + ".tmp"
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(self.persistent_items, f, indent=2, default=str)
|
||||
os.replace(temp_file, PERSISTENT_FILE)
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes PersistentQueue] Error saving persistent queue file: {e}")
|
||||
self._save_dirty = True
|
||||
if immediate:
|
||||
self._write_locked()
|
||||
else:
|
||||
self._schedule_save_locked()
|
||||
|
||||
def _schedule_save_locked(self):
|
||||
if self._save_timer is not None and self._save_timer.is_alive():
|
||||
return
|
||||
self._save_timer = threading.Timer(self.SAVE_DEBOUNCE_SECONDS, self._flush_timer)
|
||||
self._save_timer.daemon = True
|
||||
self._save_timer.start()
|
||||
|
||||
def _flush_timer(self):
|
||||
with self.lock:
|
||||
self._write_locked()
|
||||
|
||||
def _write_locked(self):
|
||||
if not self._save_dirty:
|
||||
return
|
||||
if self._read_only_newer_schema:
|
||||
# Never overwrite a queue file produced by a newer SaturnNodes.
|
||||
self._save_dirty = False
|
||||
return
|
||||
if not is_persistent_queue_enabled():
|
||||
self._save_dirty = False
|
||||
return
|
||||
try:
|
||||
os.makedirs(USER_DIR, exist_ok=True)
|
||||
temp_file = PERSISTENT_FILE + ".tmp"
|
||||
payload = {
|
||||
"schema_version": PERSISTENT_QUEUE_SCHEMA_VERSION,
|
||||
"items": self.persistent_items,
|
||||
}
|
||||
with open(temp_file, "w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, separators=(",", ":"), default=str)
|
||||
f.flush()
|
||||
try:
|
||||
os.fsync(f.fileno())
|
||||
except OSError:
|
||||
pass
|
||||
# Windows can transiently lock the destination (antivirus, indexer,
|
||||
# backup). Retry briefly instead of dropping the save.
|
||||
for attempt in range(3):
|
||||
try:
|
||||
os.replace(temp_file, PERSISTENT_FILE)
|
||||
break
|
||||
except PermissionError:
|
||||
if attempt == 2:
|
||||
raise
|
||||
time.sleep(0.05)
|
||||
self._save_dirty = False
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes PersistentQueue] Error saving persistent queue file: {e}")
|
||||
|
||||
def flush(self):
|
||||
"""Force any pending debounced write (used on shutdown)."""
|
||||
with self.lock:
|
||||
self._write_locked()
|
||||
|
||||
def add_item(self, item_tuple):
|
||||
if self.is_restoring or not is_persistent_queue_enabled():
|
||||
@@ -297,10 +470,17 @@ class PersistentQueueManager:
|
||||
if pid_str in self.batch_meta:
|
||||
entry["batch_info"] = self.batch_meta[pid_str]
|
||||
with self.lock:
|
||||
self.recently_finished.pop(pid_str, None)
|
||||
self.persistent_items = [x for x in self.persistent_items if str(x.get("prompt_id")) != pid_str]
|
||||
self.persistent_items.append(entry)
|
||||
self.save_to_file()
|
||||
print(f"[SaturnNodes PersistentQueue] Persisted prompt {pid_str} to disk.")
|
||||
if self._read_only_newer_schema:
|
||||
print(
|
||||
f"[SaturnNodes PersistentQueue] Queue file uses a newer schema; "
|
||||
f"prompt {pid_str} kept in memory only."
|
||||
)
|
||||
else:
|
||||
print(f"[SaturnNodes PersistentQueue] Persisted prompt {pid_str} to disk.")
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes PersistentQueue] Error formatting queue item for persistence: {e}")
|
||||
|
||||
@@ -312,6 +492,8 @@ class PersistentQueueManager:
|
||||
initial_count = len(self.persistent_items)
|
||||
self.persistent_items = [x for x in self.persistent_items if str(x.get("prompt_id")) != pid_str]
|
||||
self.batch_meta.pop(pid_str, None)
|
||||
self.batch_meta_updated.pop(pid_str, None)
|
||||
self._mark_finished_locked(pid_str)
|
||||
changed = len(self.persistent_items) != initial_count
|
||||
if changed and is_persistent_queue_enabled():
|
||||
self.save_to_file()
|
||||
@@ -321,8 +503,9 @@ class PersistentQueueManager:
|
||||
with self.lock:
|
||||
self.persistent_items = []
|
||||
self.batch_meta = {}
|
||||
self.batch_meta_updated = {}
|
||||
if is_persistent_queue_enabled():
|
||||
self.save_to_file()
|
||||
self.save_to_file(immediate=True)
|
||||
print("[SaturnNodes PersistentQueue] Cleared all saved queue items.")
|
||||
|
||||
def patch_server(self):
|
||||
@@ -431,6 +614,11 @@ class PersistentQueueManager:
|
||||
pause_manager.set_pause(False)
|
||||
|
||||
print(f"[SaturnNodes PersistentQueue] Restoring {len(self.persistent_items)} saved queue item(s) to client '{active_client_id}'...")
|
||||
try:
|
||||
max_retries = int(get_env_setting("PERSISTENT_QUEUE_MAX_RETRIES", "0"))
|
||||
except (TypeError, ValueError):
|
||||
max_retries = 0
|
||||
max_retries = max(0, min(10, max_retries))
|
||||
self.is_restoring = True
|
||||
max_number = 0
|
||||
restored_count = 0
|
||||
@@ -461,9 +649,19 @@ class PersistentQueueManager:
|
||||
max_number = max(max_number, int(item_list[0]))
|
||||
|
||||
server.prompt_queue.put(tuple(item_list))
|
||||
entry["retries"] = 0
|
||||
restored_count += 1
|
||||
except Exception as e:
|
||||
print(f"[SaturnNodes PersistentQueue] Error restoring prompt {entry.get('prompt_id')}: {e}")
|
||||
pid = entry.get("prompt_id")
|
||||
retries = int(entry.get("retries", 0)) + 1
|
||||
if retries > max_retries:
|
||||
print(f"[SaturnNodes PersistentQueue] Error restoring prompt {pid}: {e}. Dropping after {retries} attempt(s).")
|
||||
if pid:
|
||||
self.remove_item(pid)
|
||||
else:
|
||||
print(f"[SaturnNodes PersistentQueue] Error restoring prompt {pid}: {e}. Retry {retries}/{max_retries} on next launch.")
|
||||
entry["retries"] = retries
|
||||
self.save_to_file(immediate=True)
|
||||
finally:
|
||||
self.is_restoring = False
|
||||
self.has_claimed_once = True
|
||||
@@ -479,7 +677,7 @@ class PersistentQueueManager:
|
||||
if isinstance(entry.get("item"), list) and len(entry["item"]) > 3:
|
||||
if isinstance(entry["item"][3], dict):
|
||||
entry["item"][3]["client_id"] = active_client_id
|
||||
self.save_to_file()
|
||||
self.save_to_file(immediate=True)
|
||||
|
||||
# Broadcast updated queue state to client UI
|
||||
try:
|
||||
@@ -670,12 +868,14 @@ def _consume_power_ticket(ticket, expected_action=None):
|
||||
pause_manager = PauseQueueManager()
|
||||
persistent_manager = PersistentQueueManager()
|
||||
persistent_queue_manager = persistent_manager
|
||||
atexit.register(persistent_manager.flush)
|
||||
tray_manager = TrayIconManager(pause_manager)
|
||||
power_manager = PowerControlManager(pause_manager)
|
||||
|
||||
def setup_queue_control_routes(server):
|
||||
pause_manager.patch_all()
|
||||
persistent_manager.patch_server()
|
||||
patch_generation_times(server)
|
||||
power_manager.start_watcher()
|
||||
|
||||
if is_tray_icon_enabled():
|
||||
@@ -805,8 +1005,17 @@ def setup_queue_control_routes(server):
|
||||
except Exception:
|
||||
data = {}
|
||||
limit = data.get("limit")
|
||||
if limit is not None:
|
||||
try:
|
||||
limit = max(1, min(1000, int(limit)))
|
||||
except (TypeError, ValueError):
|
||||
limit = None
|
||||
force = bool(data.get("force", False))
|
||||
count = assets_restore_manager.restore_on_launch(server, limit=limit, force=force)
|
||||
# Scanning the output tree and parsing metadata is blocking disk I/O;
|
||||
# keep the aiohttp event loop responsive for the Restore Now button.
|
||||
count = await asyncio.to_thread(
|
||||
assets_restore_manager.restore_on_launch, server, limit=limit, force=force
|
||||
)
|
||||
return web.json_response({"success": True, "restored": count, "debug": assets_restore_manager.last_debug_report})
|
||||
|
||||
@routes.get("/saturnnodes/assets/debug")
|
||||
|
||||
+1
-12
@@ -3,22 +3,11 @@ import sys
|
||||
import threading
|
||||
import webbrowser
|
||||
from PIL import Image, ImageDraw
|
||||
from .utils import get_saturnnodes_user_dir
|
||||
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 get_env_setting(key, default_val):
|
||||
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(f"{key}="):
|
||||
return line.strip().split("=", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
return default_val
|
||||
|
||||
def is_tray_icon_enabled():
|
||||
val = get_env_setting("ENABLE_TRAY_ICON", "false").lower()
|
||||
if val not in ["true", "1", "yes"]:
|
||||
|
||||
+40
-9
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import threading
|
||||
from urllib.parse import urlparse
|
||||
import folder_paths
|
||||
|
||||
@@ -216,21 +217,51 @@ def sanitize_image_loader_folder(folder_input, default_to_output=True):
|
||||
|
||||
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")
|
||||
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(f"{key}="):
|
||||
return line.strip().split("=", 1)[1].strip()
|
||||
except Exception:
|
||||
pass
|
||||
return default_val
|
||||
|
||||
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 ]"]:
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "saturnnodes"
|
||||
description = "Essential workflow controls, mid-generation pausing, real-time queue persistence, folder watcher, and visual model loaders for ComfyUI."
|
||||
version = "2.6.1"
|
||||
version = "2.6.2"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -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,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()
|
||||
+143
-1
@@ -3,12 +3,14 @@ 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
|
||||
from nodes.queue_control import PersistentQueueManager, PauseQueueManager
|
||||
import nodes.queue_control as queue_control
|
||||
|
||||
# 36 Colors matching web/batch_queue.js
|
||||
BATCH_COLORS = [
|
||||
@@ -162,6 +164,7 @@ class TestPersistentQueueBatchMetadata(unittest.TestCase):
|
||||
# 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]
|
||||
@@ -179,6 +182,7 @@ class TestPersistentQueueBatchMetadata(unittest.TestCase):
|
||||
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)
|
||||
|
||||
@@ -191,5 +195,143 @@ class TestPersistentQueueBatchMetadata(unittest.TestCase):
|
||||
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,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,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,41 @@
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import unittest
|
||||
|
||||
|
||||
DOM_TESTS = [
|
||||
"test_power_popup_dom.mjs",
|
||||
"test_copy_prompt_dom.mjs",
|
||||
"test_floating_dock_dom.mjs",
|
||||
"test_widget_guards_dom.mjs",
|
||||
"test_queue_height_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()
|
||||
@@ -1,10 +1,20 @@
|
||||
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)
|
||||
@@ -75,6 +85,11 @@ if "folder_paths" not in sys.modules:
|
||||
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()
|
||||
|
||||
@@ -1,11 +1,18 @@
|
||||
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):
|
||||
@@ -158,5 +165,81 @@ 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,170 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
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
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
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_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,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,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");
|
||||
+192
-16
@@ -4,12 +4,21 @@ 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
|
||||
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):
|
||||
@@ -67,7 +76,7 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
patch("nodes.prompt_iterator.PromptServer.instance.send_sync"):
|
||||
|
||||
for expected_idx in range(10):
|
||||
p, rem_text, rem_count = self.node.process_queue(
|
||||
p, rem_text, rem_count, exhausted = self.node.process_queue(
|
||||
pop_mode="Sequential (Loop on End)",
|
||||
separator=">1 Empty Line",
|
||||
text=prompts_10,
|
||||
@@ -96,21 +105,21 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
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")
|
||||
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")
|
||||
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")
|
||||
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):
|
||||
@@ -137,7 +146,7 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
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")
|
||||
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)
|
||||
|
||||
@@ -160,8 +169,8 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
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")
|
||||
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")
|
||||
|
||||
@@ -171,7 +180,7 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
state_db[k]["index"] = 0
|
||||
|
||||
# Next run must be P0 again!
|
||||
p0_again, _, _ = self.node.process_queue(text=pack, unique_id="reset_test")
|
||||
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):
|
||||
@@ -195,18 +204,18 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
|
||||
# 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")
|
||||
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")
|
||||
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")
|
||||
p1, _, _, _ = self.node.process_queue(text=pack, unique_id="recover_node")
|
||||
self.assertEqual(p1, "Prompt 1")
|
||||
|
||||
def test_per_node_hash_pruning(self):
|
||||
@@ -255,7 +264,7 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
|
||||
picked = []
|
||||
for _ in range(5):
|
||||
p, _, _ = self.node.process_queue(
|
||||
p, _, _, _ = self.node.process_queue(
|
||||
pop_mode="Random (Cycle)",
|
||||
text=pack,
|
||||
unique_id="cycle_test_node"
|
||||
@@ -266,5 +275,172 @@ class TestPromptQueueIterator(unittest.TestCase):
|
||||
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");
|
||||
+48
-1
@@ -1,6 +1,10 @@
|
||||
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
|
||||
@@ -15,10 +19,22 @@ from nodes.utils import (
|
||||
sanitize_image_loader_folder,
|
||||
get_csrf_token,
|
||||
is_same_origin_or_direct,
|
||||
is_authenticated_local_request
|
||||
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):
|
||||
@@ -204,5 +220,36 @@ class TestUtils(unittest.TestCase):
|
||||
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");
|
||||
+121
-39
@@ -182,6 +182,12 @@ function renderBatchLineOnElement(element, segmentType, color) {
|
||||
element.appendChild(wrap);
|
||||
}
|
||||
|
||||
// Avoid rewriting the SVG on every status/executing refresh when the
|
||||
// visible segment and batch color have not changed.
|
||||
const signature = `${segmentType}|${color}`;
|
||||
if (wrap.dataset.saturnnodesBatchSignature === signature && wrap.innerHTML) return;
|
||||
wrap.dataset.saturnnodesBatchSignature = signature;
|
||||
|
||||
const dPath = buildSvgPath(segmentType);
|
||||
wrap.innerHTML = `
|
||||
<svg class="saturnnodes-batch-line-svg" viewBox="0 0 10 48" preserveAspectRatio="none">
|
||||
@@ -337,6 +343,29 @@ async function loadBatchFromServer() {
|
||||
} 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()) {
|
||||
@@ -348,14 +377,37 @@ async function updateQueueBatchVisuals() {
|
||||
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
|
||||
const visiblePromptIds = elements.map(el => el.getAttribute("data-job-id")).filter(Boolean);
|
||||
|
||||
// Build map of promptId -> element for visible items
|
||||
// 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) elMap.set(pid, el);
|
||||
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:
|
||||
@@ -376,11 +428,8 @@ async function updateQueueBatchVisuals() {
|
||||
const batchId = info.batchId;
|
||||
const color = info.color || BATCH_COLORS[0];
|
||||
|
||||
// Find all visible items of this batch in top-to-bottom visual order
|
||||
const occurrencesOnScreen = visiblePromptIds.filter(p => {
|
||||
const b = batchRegistry.get(p);
|
||||
return b && b.batchId === batchId;
|
||||
});
|
||||
// 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
|
||||
@@ -449,11 +498,17 @@ function setupQueueHooks() {
|
||||
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);
|
||||
}
|
||||
const itemIndex = activeBatchContext.prompts.length;
|
||||
|
||||
// Batch Graph Snapshot Guard:
|
||||
// When queueing a batch with batchCount > 1, snapshot on iteration 0 in the background
|
||||
// so canvas edits (text, LoRA selection, widgets) don't bleed into subsequent items.
|
||||
if (isSnapshotGuardEnabled() && activeBatchContext && activeBatchContext.batchCount > 1) {
|
||||
const itemIndex = activeBatchContext.prompts.length;
|
||||
if (itemIndex === 0) {
|
||||
try {
|
||||
activeBatchContext.snapshotPrompt = JSON.parse(JSON.stringify(prompt));
|
||||
@@ -465,14 +520,17 @@ function setupQueueHooks() {
|
||||
}
|
||||
}
|
||||
|
||||
const res = await origApiQueuePrompt(number, promptToSend, targets);
|
||||
let res;
|
||||
try {
|
||||
res = await origApiQueuePrompt(number, promptToSend, targets);
|
||||
} catch (e) {
|
||||
if (activeBatchContext && activeBatchContext.prompts.length === 0) {
|
||||
activeBatchContext = null;
|
||||
}
|
||||
throw e;
|
||||
}
|
||||
if (res && res.prompt_id) {
|
||||
const pid = String(res.prompt_id);
|
||||
if (!activeBatchContext) {
|
||||
activeBatchContext = getNextBatch(1);
|
||||
}
|
||||
|
||||
const itemIndex = activeBatchContext.prompts.length;
|
||||
const batchInfo = {
|
||||
prompt_id: pid,
|
||||
batchId: activeBatchContext.batchId,
|
||||
@@ -495,7 +553,7 @@ function setupQueueHooks() {
|
||||
activeBatchContext = null;
|
||||
}
|
||||
|
||||
setTimeout(updateQueueBatchVisuals, 100);
|
||||
scheduleBatchVisualRefresh(600);
|
||||
}
|
||||
return res;
|
||||
};
|
||||
@@ -507,32 +565,52 @@ app.registerExtension({
|
||||
async setup() {
|
||||
injectBatchStyles();
|
||||
setupQueueHooks();
|
||||
loadBatchFromServer();
|
||||
loadBatchFromServer().then(() => scheduleBatchVisualRefresh(1200));
|
||||
|
||||
// Listen for queue changes
|
||||
if (api && api.addEventListener) {
|
||||
api.addEventListener("status", () => {
|
||||
setTimeout(updateQueueBatchVisuals, 100);
|
||||
});
|
||||
api.addEventListener("execution_start", () => {
|
||||
setTimeout(updateQueueBatchVisuals, 80);
|
||||
});
|
||||
api.addEventListener("executing", () => {
|
||||
setTimeout(updateQueueBatchVisuals, 150);
|
||||
});
|
||||
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
|
||||
// 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 (m.target && m.target.nodeType === 1) {
|
||||
if (m.target.closest?.('[data-testid="job-assets-list"], .comfy-table') ||
|
||||
m.target.matches?.('[data-testid="job-assets-list"], [data-job-id]')) {
|
||||
shouldUpdate = true;
|
||||
break;
|
||||
}
|
||||
if (mutationTouchesQueue(m)) {
|
||||
shouldUpdate = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (shouldUpdate) {
|
||||
@@ -541,11 +619,15 @@ app.registerExtension({
|
||||
}
|
||||
});
|
||||
|
||||
observer.observe(document.body, { childList: true, subtree: true });
|
||||
observer.observe(document.body, {
|
||||
childList: true,
|
||||
subtree: true,
|
||||
attributes: true,
|
||||
attributeFilter: ["data-job-id"],
|
||||
});
|
||||
|
||||
// Initial check after load
|
||||
setTimeout(updateQueueBatchVisuals, 800);
|
||||
setTimeout(updateQueueBatchVisuals, 2000);
|
||||
// Initial checks after load, including the async /jobs hydration.
|
||||
scheduleBatchVisualRefresh(2500);
|
||||
}
|
||||
});
|
||||
|
||||
|
||||
+55
-21
@@ -1035,8 +1035,13 @@ class FloatingLatentDock {
|
||||
// Safety clamp: bottomGap must never push dock beyond upper half of screen
|
||||
bottomGap = Math.max(54, Math.min(bottomGap, Math.round(window.innerHeight * 0.75)));
|
||||
|
||||
this.dockEl.style.bottom = `${bottomGap}px`;
|
||||
this.dockEl.style.right = "8px";
|
||||
const nextBottom = `${bottomGap}px`;
|
||||
if (this.dockEl.style.bottom !== nextBottom) {
|
||||
this.dockEl.style.bottom = nextBottom;
|
||||
}
|
||||
if (this.dockEl.style.right !== "8px") {
|
||||
this.dockEl.style.right = "8px";
|
||||
}
|
||||
}
|
||||
|
||||
requestPositionUpdate() {
|
||||
@@ -1053,38 +1058,67 @@ class FloatingLatentDock {
|
||||
// Track window resize (debounced via rAF)
|
||||
window.addEventListener("resize", () => this.requestPositionUpdate(), { passive: true });
|
||||
|
||||
// Observe DOM structural changes only (minimap toggle / toolbar mount) without watching every subtree attribute
|
||||
const observer = new MutationObserver(() => {
|
||||
this.requestPositionUpdate();
|
||||
// The minimap and canvas toolbar mount/unmount deep inside the app
|
||||
// shell. A shallow body observer never sees those mutations, so the
|
||||
// dock kept a stale bottom offset and could overlap the minimap while
|
||||
// idle. Inspect added/removed subtrees for layout-relevant nodes.
|
||||
const LAYOUT_SELECTOR =
|
||||
".minimap-main-container, .litegraph-minimap, [role='toolbar'], " +
|
||||
".graph-canvas-menu, .canvas-toolbar, .side-bar";
|
||||
const isLayoutNode = (node) => {
|
||||
if (!node || node.nodeType !== 1) return false;
|
||||
if (node.matches?.(LAYOUT_SELECTOR)) return true;
|
||||
return !!node.querySelector?.(LAYOUT_SELECTOR);
|
||||
};
|
||||
|
||||
const observer = new MutationObserver((mutations) => {
|
||||
for (const mutation of mutations) {
|
||||
if (mutation.type !== "childList") continue;
|
||||
for (const node of mutation.addedNodes || []) {
|
||||
if (isLayoutNode(node)) {
|
||||
this.requestPositionUpdate();
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (const node of mutation.removedNodes || []) {
|
||||
if (isLayoutNode(node)) {
|
||||
this.requestPositionUpdate();
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
observer.observe(document.body, {
|
||||
childList: true,
|
||||
subtree: false
|
||||
subtree: true
|
||||
});
|
||||
}
|
||||
|
||||
injectToolbarButtonLoop() {
|
||||
const tryInject = () => {
|
||||
if (this.btnEl && document.body.contains(this.btnEl)) return;
|
||||
if (!this.btnEl || !document.body.contains(this.btnEl)) {
|
||||
const toolbar = document.querySelector(
|
||||
'[role="toolbar"][aria-label*="canvasToolbar" i], [role="toolbar"][aria-label*="Canvas Toolbar" i], .graph-canvas-menu, .canvas-toolbar'
|
||||
);
|
||||
if (toolbar) {
|
||||
this.createToolbarButton();
|
||||
|
||||
const toolbar = document.querySelector(
|
||||
'[role="toolbar"][aria-label*="canvasToolbar" i], [role="toolbar"][aria-label*="Canvas Toolbar" i], .graph-canvas-menu, .canvas-toolbar'
|
||||
);
|
||||
if (!toolbar) return;
|
||||
// Find Minimap button or FitView button in the toolbar to place alongside it
|
||||
const minimapBtn = toolbar.querySelector('button[aria-label*="minimap" i], button[aria-label*="Minimap" i]');
|
||||
if (minimapBtn && minimapBtn.parentElement === toolbar) {
|
||||
toolbar.insertBefore(this.btnEl, minimapBtn.nextSibling);
|
||||
} else {
|
||||
toolbar.appendChild(this.btnEl);
|
||||
}
|
||||
|
||||
this.createToolbarButton();
|
||||
|
||||
// Find Minimap button or FitView button in the toolbar to place alongside it
|
||||
const minimapBtn = toolbar.querySelector('button[aria-label*="minimap" i], button[aria-label*="Minimap" i]');
|
||||
if (minimapBtn && minimapBtn.parentElement === toolbar) {
|
||||
toolbar.insertBefore(this.btnEl, minimapBtn.nextSibling);
|
||||
} else {
|
||||
toolbar.appendChild(this.btnEl);
|
||||
this.updateVisibility();
|
||||
}
|
||||
}
|
||||
|
||||
this.updateVisibility();
|
||||
this.updatePosition();
|
||||
// Catch layout-only toggles (class/style changes) that produce no
|
||||
// childList mutation. rAF-debounced and a no-op while hidden.
|
||||
this.requestPositionUpdate();
|
||||
};
|
||||
|
||||
tryInject();
|
||||
|
||||
+61
-12
@@ -358,6 +358,20 @@ app.registerExtension({
|
||||
|
||||
// 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;
|
||||
@@ -370,22 +384,41 @@ function isQueueItemElement(el) {
|
||||
|
||||
function registerAssetInteraction(target) {
|
||||
if (!target) return;
|
||||
const card = target.closest?.(
|
||||
"div[data-virtual-grid-item], [data-asset-id], .asset-card, [data-testid='asset-card'], [data-node-id], .lg-node, .comfy-image-preview, .group"
|
||||
);
|
||||
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
|
||||
if (lastInteractedAssetCard) {
|
||||
const img = lastInteractedAssetCard.querySelector("img");
|
||||
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
|
||||
@@ -464,7 +497,11 @@ function injectCopyPromptNextToDownload(downloadBtn) {
|
||||
copyBtn.onclick = async (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
const activeImg = card.querySelector("img") || img;
|
||||
// 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
|
||||
@@ -503,7 +540,8 @@ function injectCopyPromptNextToDownload(downloadBtn) {
|
||||
inspectBtn.onclick = (e) => {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
const target = card.querySelector(".aspect-square, [class*='aspect-square'], img") || img;
|
||||
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 }));
|
||||
}
|
||||
@@ -609,7 +647,7 @@ function injectContextMenuCopy(contextMenu) {
|
||||
|
||||
// 5. Fallback Hook for Older Frontend Action Bars
|
||||
function injectHoverCopyAction(overlayBar) {
|
||||
if (!overlayBar || overlayBar.querySelector(".saturnnodes-hover-copy")) return;
|
||||
if (!overlayBar) return;
|
||||
if (!isCopyEnabled("SaturnNodes.📋 Prompt Actions.01_EnableAssetsCopyPromptButton")) return;
|
||||
|
||||
if (isQueueItemElement(overlayBar)) return;
|
||||
@@ -626,9 +664,7 @@ function injectHoverCopyAction(overlayBar) {
|
||||
overlayBar.querySelector('button[aria-label*="zoom" i]')?.parentElement ||
|
||||
overlayBar;
|
||||
|
||||
const parentCard = overlayBar.closest(
|
||||
"div[data-virtual-grid-item], .asset-card, [data-testid='asset-card'], [data-node-id], .lg-node, .comfy-image-preview"
|
||||
);
|
||||
const parentCard = overlayBar.closest(ASSET_CARD_SELECTOR);
|
||||
if (!parentCard) return;
|
||||
if (isQueueItemElement(parentCard)) return;
|
||||
|
||||
@@ -636,18 +672,31 @@ function injectHoverCopyAction(overlayBar) {
|
||||
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();
|
||||
const success = await copyImagePrompt(img.src);
|
||||
// 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);
|
||||
|
||||
@@ -2,6 +2,7 @@ import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
import { isSaturnColorsEnabled } from "../saturnnodes_colors.js";
|
||||
import { authenticatedFetch } from "./auth_helper.js";
|
||||
import { resolveLiveWidgetElement } from "./live_dom.js";
|
||||
|
||||
// Hash folder name to get a consistent color for the border
|
||||
function getFolderBorderColor(folderName) {
|
||||
@@ -1270,10 +1271,18 @@ app.registerExtension({
|
||||
|
||||
const updateLoraBadgeColor = () => {
|
||||
const enabled = isSaturnColorsEnabled();
|
||||
if (node._loraBadgeSpan) {
|
||||
node._loraBadgeSpan.style.color = enabled ? "#f59e0b" : "var(--text-muted, #a1a1aa)";
|
||||
const badgeEl = resolveLiveWidgetElement(
|
||||
node._loraBadgeEl,
|
||||
node.id,
|
||||
"lora_count_preview"
|
||||
);
|
||||
if (badgeEl) node._loraBadgeEl = badgeEl;
|
||||
const badgeSpan = badgeEl?.querySelector?.("span") || node._loraBadgeSpan;
|
||||
if (badgeSpan) node._loraBadgeSpan = badgeSpan;
|
||||
|
||||
if (badgeSpan) {
|
||||
badgeSpan.style.color = enabled ? "#f59e0b" : "var(--text-muted, #a1a1aa)";
|
||||
}
|
||||
const badgeEl = node._loraBadgeEl;
|
||||
if (badgeEl) {
|
||||
badgeEl.style.background = enabled
|
||||
? "rgba(180, 140, 95, 0.12)"
|
||||
@@ -1314,8 +1323,16 @@ app.registerExtension({
|
||||
const count = getSelectedLoras().length;
|
||||
node._selectedLoraCount = count;
|
||||
const label = `${count} LoRA${count === 1 ? "" : "s"} Selected`;
|
||||
if (node._loraBadgeSpan) {
|
||||
node._loraBadgeSpan.textContent = label;
|
||||
const liveBadgeEl = resolveLiveWidgetElement(
|
||||
node._loraBadgeEl,
|
||||
node.id,
|
||||
"lora_count_preview"
|
||||
);
|
||||
if (liveBadgeEl) node._loraBadgeEl = liveBadgeEl;
|
||||
const liveBadgeSpan = liveBadgeEl?.querySelector?.("span") || node._loraBadgeSpan;
|
||||
if (liveBadgeSpan) {
|
||||
node._loraBadgeSpan = liveBadgeSpan;
|
||||
liveBadgeSpan.textContent = label;
|
||||
}
|
||||
if (loraCountWidget) {
|
||||
loraCountWidget.value = label;
|
||||
|
||||
@@ -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}"]`
|
||||
);
|
||||
});
|
||||
}
|
||||
+23
-5
@@ -12,6 +12,7 @@ let currentPowerState = {
|
||||
};
|
||||
|
||||
let activePopup = null;
|
||||
let activePopupCleanup = null;
|
||||
|
||||
function injectStyles() {
|
||||
if (document.getElementById(STYLE_ID)) return;
|
||||
@@ -296,6 +297,15 @@ function checkProcessManagementEnabled() {
|
||||
}
|
||||
|
||||
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;
|
||||
@@ -307,10 +317,14 @@ async function powerFetch(endpointSuffix, options = {}) {
|
||||
}
|
||||
|
||||
function showPowerSidePopup(anchorElement) {
|
||||
if (activePopup) {
|
||||
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";
|
||||
@@ -607,18 +621,22 @@ function showPowerSidePopup(anchorElement) {
|
||||
const outsideListener = (e) => {
|
||||
if (popup && !popup.contains(e.target) && !anchorElement.contains(e.target)) {
|
||||
closePowerPopup();
|
||||
document.removeEventListener("pointerdown", outsideListener);
|
||||
document.removeEventListener("keydown", keyListener);
|
||||
}
|
||||
};
|
||||
const keyListener = (e) => {
|
||||
if (e.key === "Escape") {
|
||||
closePowerPopup();
|
||||
document.removeEventListener("pointerdown", outsideListener);
|
||||
document.removeEventListener("keydown", keyListener);
|
||||
}
|
||||
};
|
||||
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);
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
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;
|
||||
@@ -104,10 +105,18 @@ app.registerExtension({
|
||||
|
||||
function updatePromptBadgeColor() {
|
||||
const enabled = isSaturnColorsEnabled();
|
||||
if (node._promptBadgeSpan) {
|
||||
node._promptBadgeSpan.style.color = enabled ? "#f59e0b" : "var(--text-muted, #a1a1aa)";
|
||||
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)";
|
||||
}
|
||||
const badgeEl = node._promptBadgeEl;
|
||||
if (badgeEl) {
|
||||
badgeEl.style.background = enabled
|
||||
? "rgba(180, 140, 95, 0.12)"
|
||||
@@ -126,8 +135,16 @@ app.registerExtension({
|
||||
const count = countPromptsInText(text, sep, regex);
|
||||
node._promptCount = count;
|
||||
const label = `${count} Prompt${count === 1 ? "" : "s"}`;
|
||||
if (node._promptBadgeSpan) {
|
||||
node._promptBadgeSpan.textContent = label;
|
||||
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;
|
||||
|
||||
+23
-3
@@ -1,6 +1,7 @@
|
||||
import { app } from "/scripts/app.js";
|
||||
import { api } from "/scripts/api.js";
|
||||
import { PreviewManager } from "./js/preview_manager.js";
|
||||
import { resolveLiveElement } from "./js/live_dom.js";
|
||||
|
||||
// Global CSS injection to guarantee zero-expansion in Vue Node 2.0 and LiteGraph
|
||||
if (typeof document !== "undefined") {
|
||||
@@ -370,9 +371,28 @@ app.registerExtension({
|
||||
}
|
||||
|
||||
if (this.previewImgElement && this.previewFallbackText && img && img.src) {
|
||||
this.previewFallbackText.style.display = "none";
|
||||
this.previewImgElement.src = img.src;
|
||||
this.previewImgElement.style.display = "block";
|
||||
// The Vue node renderer can remount the DOM widget while the
|
||||
// node object keeps its cached element references. Resolve the
|
||||
// element that is actually mounted for this node before writing.
|
||||
const liveImg = resolveLiveElement(this.previewImgElement, () => {
|
||||
if (this.id === undefined || this.id === null) return null;
|
||||
return document.querySelector(
|
||||
`[data-node-id="${this.id}"] .saturnnodes-live-preview-img, ` +
|
||||
`[data-widgets-grid-node-id="${this.id}"] .saturnnodes-live-preview-img`
|
||||
);
|
||||
});
|
||||
if (liveImg) this.previewImgElement = liveImg;
|
||||
|
||||
const liveFallback = resolveLiveElement(this.previewFallbackText, () =>
|
||||
liveImg?.parentElement?.querySelector?.(".saturnnodes-live-preview-fallback")
|
||||
);
|
||||
if (liveFallback) this.previewFallbackText = liveFallback;
|
||||
|
||||
if (liveFallback) liveFallback.style.display = "none";
|
||||
if (liveImg) {
|
||||
liveImg.src = img.src;
|
||||
liveImg.style.display = "block";
|
||||
}
|
||||
}
|
||||
|
||||
if (app.graph && typeof app.graph.setDirtyCanvas === "function") {
|
||||
|
||||
+189
-4
@@ -15,6 +15,55 @@ async function postSaturnNodesSettings(bodyObj) {
|
||||
}
|
||||
}
|
||||
|
||||
let _backendSettings = null;
|
||||
const BACKEND_SETTING_MAP = {
|
||||
"SaturnNodes.🖼️ Visual Loaders.01_CivitaiApiKey": "civitai_api_key",
|
||||
"SaturnNodes.🖼️ Visual Loaders.02_EnableCivitaiScraping": "enable_civitai_scraping",
|
||||
"SaturnNodes.🖼️ Visual Loaders.03_TMDBApiKey": "tmdb_api_key",
|
||||
"SaturnNodes.🖼️ Visual Loaders.04_EnableTMDBScraping": "enable_tmdb_scraping",
|
||||
"SaturnNodes.🖼️ Visual Loaders.05_EnableLoraUsage": "enable_lora_usage",
|
||||
"SaturnNodes.🖼️ Visual Loaders.08_ImageMetadataScanLimit": "image_loader_metadata_limit",
|
||||
"SaturnNodes.⏸️ Pause Controls.01_DefaultStateOnLaunch": "default_pause_state",
|
||||
"SaturnNodes.⏸️ Pause Controls.02_DefaultPauseAction": "default_pause_mode",
|
||||
"SaturnNodes.⏸️ Pause Controls.02b_PersistPauseState": "persist_pause_state",
|
||||
"SaturnNodes.⏸️ Pause Controls.06_EnableTrayIcon": "enable_tray_icon",
|
||||
"SaturnNodes.⏸️ Pause Controls.07_AllowProcessManagement": "enable_process_management",
|
||||
"SaturnNodes.🔄 Prompt Iterator.01_ClearOnLaunch": "clear_prompt_iterator_on_launch",
|
||||
"SaturnNodes.💾 Persistent Queue.01_EnablePersistentQueue": "enable_persistent_queue",
|
||||
"SaturnNodes.💾 Persistent Queue.02_RecoveryLaunchState": "persistent_queue_restored_state",
|
||||
"SaturnNodes.💾 Persistent Queue.03_MaxRestoreRetries": "persistent_queue_max_retries",
|
||||
"SaturnNodes.🖼️ Assets Restore.01_RestoreAssetsOnLaunch": "enable_assets_restore",
|
||||
"SaturnNodes.🖼️ Assets Restore.02_RestoredAssetsCount": "restore_assets_count",
|
||||
"SaturnNodes.🖼️ Assets Restore.03_RestoreAssetsRecursive": "restore_assets_recursive",
|
||||
"SaturnNodes.🖼️ Assets Restore.04_ShowRestoredGenerationTime": "show_restored_generation_time",
|
||||
"SaturnNodes.🛡️ Security.01_AllowLocalFileExecution": "enable_local_file_execution"
|
||||
};
|
||||
|
||||
async function fetchBackendSettings() {
|
||||
try {
|
||||
const res = await authenticatedFetch("/saturnnodes/settings");
|
||||
if (res && res.ok) {
|
||||
_backendSettings = await res.json();
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn("[SaturnNodes Settings] Failed to load backend .env settings:", e);
|
||||
}
|
||||
}
|
||||
|
||||
function getBackendInitialSetting(saturnId) {
|
||||
if (!_backendSettings) return undefined;
|
||||
const backendKey = BACKEND_SETTING_MAP[saturnId];
|
||||
if (!backendKey || !(backendKey in _backendSettings)) return undefined;
|
||||
const value = _backendSettings[backendKey];
|
||||
if (backendKey === "default_pause_state") {
|
||||
return String(value).toLowerCase() === "paused" ? "Paused" : "Running";
|
||||
}
|
||||
if (backendKey === "default_pause_mode") {
|
||||
return String(value).toLowerCase().includes("instant") ? "Instant Resume Node" : "Finish Active Prompt";
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
let _disabledNodesCache = null;
|
||||
let _nodesInfoCache = null;
|
||||
async function fetchDisabledNodes() {
|
||||
@@ -716,7 +765,12 @@ function renderSettingButton(label, workingText, successText, onClickHandler) {
|
||||
}
|
||||
|
||||
function getInitialSetting(saturnId, defaultVal) {
|
||||
// 1. Check ComfyUI settings API (getSettingValue / get)
|
||||
// 1. The backend .env is the source of truth when it defines the value;
|
||||
// this keeps the settings UI and the server in sync after manual edits.
|
||||
const backendValue = getBackendInitialSetting(saturnId);
|
||||
if (backendValue !== undefined && backendValue !== null) return backendValue;
|
||||
|
||||
// 2. Check ComfyUI settings API (getSettingValue / get)
|
||||
if (typeof app.ui?.settings?.getSettingValue === "function") {
|
||||
try {
|
||||
const val = app.ui.settings.getSettingValue(saturnId);
|
||||
@@ -730,7 +784,7 @@ function getInitialSetting(saturnId, defaultVal) {
|
||||
} catch (_) {}
|
||||
}
|
||||
|
||||
// 2. Check localStorage
|
||||
// 3. Check localStorage
|
||||
if (typeof localStorage !== "undefined") {
|
||||
const val = localStorage.getItem(`Comfy.Settings.${saturnId}`);
|
||||
if (val !== null && val !== undefined) {
|
||||
@@ -740,9 +794,33 @@ function getInitialSetting(saturnId, defaultVal) {
|
||||
return defaultVal;
|
||||
}
|
||||
|
||||
/**
|
||||
* Re-apply the saved queue height to lists that mounted after setup.
|
||||
*
|
||||
* The legacy v1 queue can mount or recreate its `.comfy-list` independently
|
||||
* of the overlay, so `attachQueueResizeHandle()` (which early-returns once the
|
||||
* handle exists) is not enough to keep the height budget applied.
|
||||
*
|
||||
* @returns {boolean} true when the apply callback ran for a fresh list.
|
||||
*/
|
||||
function reapplyQueueHeightToMountedLists(overlay, heightPx, applyFn) {
|
||||
if (!overlay || !heightPx || heightPx <= 0 || typeof applyFn !== "function") return false;
|
||||
const lists = overlay.querySelectorAll?.(".comfy-list") || [];
|
||||
const marker = String(heightPx);
|
||||
for (const listEl of lists) {
|
||||
if (listEl.dataset?.saturnQueueHeight === marker) continue;
|
||||
applyFn(heightPx);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "ComfyUI.SaturnNodes.Settings",
|
||||
async setup() {
|
||||
// Load .env-backed settings before registering defaults so the UI and
|
||||
// backend never disagree about effective values.
|
||||
await fetchBackendSettings();
|
||||
|
||||
// =========================================================================
|
||||
// GROUP: 🖼️ Visual Loaders
|
||||
@@ -821,7 +899,7 @@ app.registerExtension({
|
||||
id: "SaturnNodes.🖼️ Visual Loaders.02_EnableCivitaiScraping",
|
||||
name: "Enable Civitai Auto-Scraping",
|
||||
type: "boolean",
|
||||
defaultValue: getInitialSetting("SaturnNodes.🖼️ Visual Loaders.02_EnableCivitaiScraping", true),
|
||||
defaultValue: getInitialSetting("SaturnNodes.🖼️ Visual Loaders.02_EnableCivitaiScraping", false),
|
||||
tooltip: "Toggles automated downloading of preview thumbnails for new LoRAs from Civitai via SHA256 file hashes. Note: SHA256 hash searching will always work regardless of this setting when matching local models.",
|
||||
onChange(value) {
|
||||
postSaturnNodesSettings({ enable_civitai_scraping: value ? "true" : "false" });
|
||||
@@ -918,6 +996,19 @@ app.registerExtension({
|
||||
}
|
||||
});
|
||||
|
||||
// 3.8 Visual Image Loader Metadata Scan Limit
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.🖼️ Visual Loaders.08_ImageMetadataScanLimit",
|
||||
name: "Image Loader Metadata Scan Limit",
|
||||
type: "number",
|
||||
defaultValue: getInitialSetting("SaturnNodes.🖼️ Visual Loaders.08_ImageMetadataScanLimit", 128),
|
||||
tooltip: "How many of the newest images in a folder are scanned for embedded prompts in the Visual Image Loader (0 disables metadata extraction, max 2000).",
|
||||
onChange(value) {
|
||||
const limit = Math.max(0, Math.min(2000, parseInt(value, 10) || 0));
|
||||
postSaturnNodesSettings({ image_loader_metadata_limit: limit });
|
||||
}
|
||||
});
|
||||
|
||||
// =========================================================================
|
||||
// GROUP: 👁️ Live Preview
|
||||
// =========================================================================
|
||||
@@ -1121,6 +1212,18 @@ app.registerExtension({
|
||||
}
|
||||
});
|
||||
|
||||
// 4.2b Persist Manual Pause State Across Restarts
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.⏸️ Pause Controls.02b_PersistPauseState",
|
||||
name: "Persist Pause State Across Restarts",
|
||||
type: "boolean",
|
||||
defaultValue: getInitialSetting("SaturnNodes.⏸️ Pause Controls.02b_PersistPauseState", false),
|
||||
tooltip: "When enabled, the last manual Pause/Continue state and pause mode are restored on the next ComfyUI start (overrides the launch defaults above).",
|
||||
onChange(value) {
|
||||
postSaturnNodesSettings({ persist_pause_state: value ? "true" : "false" });
|
||||
}
|
||||
});
|
||||
|
||||
// 4.3 Enable Pause Queue Toolbar Button
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.⏸️ Pause Controls.03_EnableToolbarButton",
|
||||
@@ -1281,6 +1384,11 @@ app.registerExtension({
|
||||
} else {
|
||||
listEl.style.removeProperty("max-height");
|
||||
}
|
||||
// Marker lets the overlay observer detect lists that mounted
|
||||
// after this pass without rewriting styles on every mutation.
|
||||
if (listEl.dataset) {
|
||||
listEl.dataset.saturnQueueHeight = String(saturnQueueHeightPx);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -1381,6 +1489,11 @@ app.registerExtension({
|
||||
const overlay = document.querySelector('[data-testid="queue-progress-overlay"]');
|
||||
if (overlay) {
|
||||
attachQueueResizeHandle(overlay);
|
||||
reapplyQueueHeightToMountedLists(
|
||||
overlay,
|
||||
saturnQueueHeightPx,
|
||||
applyQueueOverlayHeight
|
||||
);
|
||||
}
|
||||
});
|
||||
queueOverlayObserver.observe(document.body, { childList: true, subtree: true });
|
||||
@@ -1449,6 +1562,19 @@ app.registerExtension({
|
||||
}
|
||||
});
|
||||
|
||||
// 5.3 Max Restore Attempts for Failed Queue Items
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.💾 Persistent Queue.03_MaxRestoreRetries",
|
||||
name: "Max Restore Retries",
|
||||
type: "number",
|
||||
defaultValue: getInitialSetting("SaturnNodes.💾 Persistent Queue.03_MaxRestoreRetries", 0),
|
||||
tooltip: "How many times a queue item that fails to restore may be retried on later launches before it is dropped (0 = drop immediately, max 10).",
|
||||
onChange(value) {
|
||||
const retries = Math.max(0, Math.min(10, parseInt(value, 10) || 0));
|
||||
postSaturnNodesSettings({ persistent_queue_max_retries: retries });
|
||||
}
|
||||
});
|
||||
|
||||
// =========================================================================
|
||||
// GROUP: 🖼️ Assets Restore
|
||||
// =========================================================================
|
||||
@@ -1478,6 +1604,65 @@ app.registerExtension({
|
||||
}
|
||||
});
|
||||
|
||||
// 6.3 Scan Output Subfolders Recursively
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.🖼️ Assets Restore.03_RestoreAssetsRecursive",
|
||||
name: "Scan Output Subfolders Recursively",
|
||||
type: "boolean",
|
||||
defaultValue: getInitialSetting("SaturnNodes.🖼️ Assets Restore.03_RestoreAssetsRecursive", true),
|
||||
tooltip: "When enabled, restored assets include images in nested output subfolders, not just the top-level output folder.",
|
||||
onChange(value) {
|
||||
postSaturnNodesSettings({ restore_assets_recursive: value ? "true" : "false" });
|
||||
}
|
||||
});
|
||||
|
||||
// 6.4 Show Generation Time on Restored Assets
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.🖼️ Assets Restore.04_ShowRestoredGenerationTime",
|
||||
name: "Show Generation Time on Restored Assets",
|
||||
type: "boolean",
|
||||
defaultValue: getInitialSetting("SaturnNodes.🖼️ Assets Restore.04_ShowRestoredGenerationTime", true),
|
||||
tooltip: "Uses SaturnNodes' generation-time records to show real execution times on restored history entries. Older images keep their file mtime for ordering but show no duration when none was recorded.",
|
||||
onChange(value) {
|
||||
postSaturnNodesSettings({ show_restored_generation_time: value ? "true" : "false" });
|
||||
}
|
||||
});
|
||||
|
||||
// 6.5 Restore Assets Now Button
|
||||
const handleRestoreAssetsNow = async () => {
|
||||
const approved = confirm("🪐 Restore assets now?\n\nThis scans the output folder and adds the newest images to the Assets / History pane without restarting ComfyUI.");
|
||||
if (!approved) return null;
|
||||
const resp = await authenticatedFetch("/saturnnodes/assets/restore", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ force: true })
|
||||
});
|
||||
if (!resp || !resp.ok) throw new Error("Restore request failed");
|
||||
return await resp.json();
|
||||
};
|
||||
|
||||
const renderRestoreAssetsNowBtn = renderSettingButton("🖼️ Restore Now", "⏳ Restoring...", "✅ Restored!", async () => {
|
||||
const result = await handleRestoreAssetsNow();
|
||||
if (result) {
|
||||
const debug = result.debug || {};
|
||||
alert(`✅ Restored ${result.restored || 0} asset(s).\n\nScanned: ${debug.total_scanned_images ?? 0}\nElapsed: ${debug.elapsed_ms ?? 0} ms`);
|
||||
}
|
||||
});
|
||||
|
||||
app.ui.settings.addSetting({
|
||||
id: "SaturnNodes.🖼️ Assets Restore.05_RestoreNow",
|
||||
name: "Restore Assets Now",
|
||||
type: renderRestoreAssetsNowBtn,
|
||||
render: renderRestoreAssetsNowBtn,
|
||||
defaultValue: "🖼️ Restore Now",
|
||||
tooltip: "Immediately scan the output folder and restore the newest images into the Assets / History pane.",
|
||||
attrs: {
|
||||
readOnly: true,
|
||||
style: "display: flex; justify-content: flex-end; width: 100%; border: none !important; background: transparent !important; padding: 0 !important; box-shadow: none !important;",
|
||||
onClick: handleRestoreAssetsNow
|
||||
}
|
||||
});
|
||||
|
||||
// =========================================================================
|
||||
// GROUP: 🩺 Diagnostics
|
||||
// =========================================================================
|
||||
@@ -1682,7 +1867,7 @@ function setupSettingsFooterObserver() {
|
||||
<div style="display: flex; align-items: center; gap: 8px;">
|
||||
<span style="font-size: 16px;">🪐</span>
|
||||
<span style="font-weight: 600; color: #e4e4e7;">ComfyUI-SaturnNodes</span>
|
||||
<span style="font-size: 11px; padding: 2px 7px; border-radius: 4px; background: rgba(180, 140, 95, 0.18); color: #decbb2; border: 1px solid rgba(180, 140, 95, 0.35);">v2.6.1 (dev)</span>
|
||||
<span style="font-size: 11px; padding: 2px 7px; border-radius: 4px; background: rgba(180, 140, 95, 0.18); color: #decbb2; border: 1px solid rgba(180, 140, 95, 0.35);">v2.6.2</span>
|
||||
</div>
|
||||
<div style="display: flex; align-items: center; gap: 10px; flex-wrap: wrap;">
|
||||
<a href="https://github.com/KOFiblto/ComfyUI-SaturnNodes" target="_blank" rel="noopener noreferrer" style="color: #decbb2; text-decoration: underline; text-decoration-color: rgba(222, 203, 178, 0.4); text-underline-offset: 3px; font-weight: 500; font-size: 12px; display: inline-flex; align-items: center; gap: 4px; padding: 4px 10px; border-radius: 6px; background: rgba(255, 255, 255, 0.05); border: 1px solid rgba(255, 255, 255, 0.12); transition: all 0.2s ease;" onmouseover="this.style.color='#ffffff'; this.style.background='rgba(255,255,255,0.12)'; this.style.borderColor='rgba(255,255,255,0.25)';" onmouseout="this.style.color='#decbb2'; this.style.background='rgba(255,255,255,0.05)'; this.style.borderColor='rgba(255,255,255,0.12)';">
|
||||
|
||||
Reference in New Issue
Block a user