diff --git a/.coveragerc b/.coveragerc index 9eded8d..4bd6942 100644 --- a/.coveragerc +++ b/.coveragerc @@ -14,6 +14,10 @@ omit = */node_modules/* */static/* */media/* + */htmlcov/* + */.claude/* + restart_gallery.py + standalone_tagger.py branch = true parallel = true diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..a7d3b01 --- /dev/null +++ b/.flake8 @@ -0,0 +1,6 @@ +[flake8] +max-line-length = 88 +extend-ignore = E203, W503 +exclude = reference/, .git, __pycache__, web/, venv/, htmlcov/, build/, dist/, .claude/, docs/, standalone_tagger.py +per-file-ignores = + tests/*:E402 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0c27c74..9a2425a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -11,9 +11,12 @@ on: jobs: test: - runs-on: ubuntu-latest + name: py${{ matrix.python-version }} on ${{ matrix.os }} + runs-on: ${{ matrix.os }} strategy: + fail-fast: false matrix: + os: [ubuntu-latest, windows-latest, macos-latest] python-version: ["3.10", "3.11", "3.12"] steps: @@ -23,14 +26,8 @@ jobs: uses: actions/setup-python@v7 with: python-version: ${{ matrix.python-version }} - - - name: Cache pip dependencies - uses: actions/cache@v6 - with: - path: ~/.cache/pip - key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('requirements-dev.txt') }} - restore-keys: | - ${{ runner.os }}-pip-${{ matrix.python-version }}- + cache: pip + cache-dependency-path: requirements-dev.txt - name: Install dependencies run: | @@ -38,9 +35,45 @@ jobs: pip install -r requirements-dev.txt - name: Run tests - run: python -m unittest discover tests/ -v + run: python -m unittest discover -s tests -v + + coverage: + name: coverage gate (85%) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: requirements-dev.txt + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements-dev.txt + + - name: Measure coverage + run: | + python -m coverage run -m unittest discover -s tests + python -m coverage combine + python -m coverage report --fail-under=85 + + lint: + name: black and flake8 + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-python@v7 + with: + python-version: "3.12" + - run: pip install black flake8 + - run: black --check . + - run: flake8 . js-test: + name: frontend modules runs-on: ubuntu-latest steps: - uses: actions/checkout@v7 diff --git a/.gitignore b/.gitignore index 59734d7..c335b78 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,12 @@ AGENTS.md CLAUDE.md # Local files +# Local PromptManager state (config.json holds API keys and machine paths) +config.json +resume.sh +*.backup_* +promptmanager_logs/ +custom_data/ data/ docs/ logs/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1b69bc2..d7efa4e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ repos: - id: flake8 name: flake8 - entry: flake8 --select=E9,F63,F7,F82 --show-source --exclude=reference/,web/ + entry: flake8 --show-source language: system types: [python] diff --git a/README.md b/README.md index 9e9fa91..e09a25a 100644 --- a/README.md +++ b/README.md @@ -164,7 +164,7 @@ _Powered by [ViewerJS](https://github.com/fengyuanchen/viewerjs) for professiona - **Full-text search** across all stored prompts - **Category filtering** for organized prompt collections -- **Tag-based search** with support for multiple tags +- **Tag-based search** with support for multiple tags, with autocomplete that narrows to tags seen together with the ones already typed - **Rating filters** to find your best prompts - **Date range filtering** for temporal searches - **Recent prompts** quick access @@ -306,9 +306,12 @@ The comprehensive gallery system provides complete access to all your ComfyUI ou Configure your viewing experience in the settings panel: -- **Performance Settings**: Grid columns, image quality, lazy loading +- **Display Settings**: Default view, items per page, grid columns, media info, infinite scrolling +- **Performance Settings**: Image quality, lazy loading - **Video Settings**: Autoplay, mute by default, loop videos -- **Thumbnail Management**: Generate optimized thumbnails for faster loading +- **Thumbnail Management**: Generate optimized thumbnails for faster loading (runs on the configured worker threads; a run that stops at the per-run cap tells you how many remain) +- **Find Duplicate Images**: Hashes every file on the worker threads with a progress bar; close the window and it keeps running, results wait in a pill +- **Rescan Output Folder**: Drops the server-side listing cache and rewalks the folders #### **πŸ–ΌοΈ Image Viewing Experience** @@ -358,14 +361,16 @@ Use the standalone metadata viewer to analyze any ComfyUI-generated PNG: Import existing ComfyUI images into your database: 1. **Open admin dashboard** and click **"πŸ“Έ Scan Images"** -2. **Start scan** to analyze your entire output directory -3. **Monitor progress** with real-time status updates +2. **Start scan** to analyze every configured output directory +3. **Monitor progress** with real-time status updates. The scan runs on the server: closing the window only minimizes it to a progress pill, and reloading the page re-attaches to a running scan 4. **Review results** showing: - Images found and processed - Prompts extracted and linked - Any errors or issues encountered 5. **Access imported data** through the normal gallery interface +The scan reads PNG text chunks and A1111-style parameters, and, when `ffprobe` is installed, the prompt ComfyUI's video savers embed in MP4/MOV/WebM files. Files are read on the number of **worker threads** set in Settings β†’ Performance (default: half your cores), which makes a large library several times faster. + ### 🏷️ AI AutoTag Automatically tag your entire image collection using AI vision models: @@ -664,7 +669,8 @@ CREATE TABLE generated_images ( - **`utils/diagnostics.py`** - System diagnostics and health checks - **`web/admin.html`** - Advanced admin dashboard with metadata panel - **`web/index.html`** - Simple web interface -- **`web/js/prompt_manager.js`** - Dashboard JavaScript +- **`web/comfy/prompt_manager.js`** - ComfyUI canvas extension (the only file ComfyUI loads) +- **`web/js/admin.js`**, **`web/js/gallery.js`** - Dashboard and gallery JavaScript - **`web/js/tags-page.js`** - Tag management JavaScript - **`web/metadata.html`** - Standalone PNG metadata viewer @@ -699,8 +705,11 @@ ComfyUI_PromptManager/ β”‚ β”œβ”€β”€ gallery.html # Image gallery β”‚ β”œβ”€β”€ index.html # Simple web interface β”‚ β”œβ”€β”€ metadata.html # Standalone metadata viewer +β”‚ β”œβ”€β”€ comfy/ +β”‚ β”‚ └── prompt_manager.js # ComfyUI canvas extension (WEB_DIRECTORY) β”‚ └── js/ -β”‚ β”œβ”€β”€ prompt_manager.js # Dashboard JavaScript +β”‚ β”œβ”€β”€ admin.js # Dashboard JavaScript +β”‚ β”œβ”€β”€ gallery.js # Gallery JavaScript β”‚ └── tags-page.js # Tag management JavaScript β”œβ”€β”€ tests/ β”‚ β”œβ”€β”€ __init__.py @@ -763,11 +772,15 @@ Configure the web interface behavior: # Web UI settings (PromptManagerConfig class) RESULT_TIMEOUT = 5 # Seconds to auto-hide results in ComfyUI node WEBUI_DISPLAY_MODE = 'newtab' # 'popup' or 'newtab' for Web UI button +INFINITE_SCROLL = False # Dashboard list loads the next page on scroll by default SHOW_TEST_BUTTON = False # Show API test button in node UI ``` +All of these are editable in the dashboard Settings modal and saved to `config.json` (which is private to the installation and holds the CivitAI key, so it is never committed). + ### Performance Tuning +- **Worker threads** (Settings β†’ Performance): threads used by the image rescan, thumbnail generation and duplicate detection. Default is half the detected cores, maximum is the core count. Pillow releases the GIL while decoding, so the scan and thumbnails scale almost linearly; keep a few cores free while ComfyUI is generating. - The database automatically creates indexes for optimal search performance - Regular `VACUUM` operations keep the database optimized - Consider backing up the database periodically @@ -811,9 +824,14 @@ db.model.backup_database("backup_prompts.db") ```bash cd ComfyUI_PromptManager -python -m pytest tests/ -v +pip install -r requirements-dev.txt +python -m unittest discover -s tests # Python (what CI runs) +node --test tests/js/*.test.js # frontend pure modules +python -m coverage run -m unittest discover -s tests && python -m coverage combine && python -m coverage report ``` +CI runs the suite on Linux, Windows and macOS for Python 3.10 to 3.12 with an 85% coverage gate. Tests that delete a database file first check that nothing in the process still holds it open (`tests/open_handles.py`), so a connection leak fails on Linux the way it would on Windows. + ### Code Style The project follows PEP 8 guidelines with: @@ -874,6 +892,10 @@ MIT License - see LICENSE file for details. ### Recently Completed +- **βœ… Security & Reliability Hardening (3.3.0)**: validated gallery roots, contained file serving, safe backups and restores, background scans, 85%+ test coverage on all three operating systems +- **βœ… Worker Threads**: scan, thumbnails and duplicate detection run on a configurable thread pool +- **βœ… Tag Autocomplete & Infinite Scrolling**: faster browsing in the dashboard and gallery +- **βœ… Queue-Time Run Counting (3.2.4)**: run counts and image linking follow the ComfyUI queue - **βœ… LoRA Manager Integration**: Import LoRA metadata, trigger words, and preview images - **βœ… Folder Filter**: Browse and filter prompts by output subdirectory - **βœ… Multi-Directory Gallery**: Scan multiple output directories simultaneously @@ -893,6 +915,26 @@ MIT License - see LICENSE file for details. ## Changelog +### v3.3.0 (Hardening, Background Jobs & Worker Threads) + +- **πŸ”’ Security**: gallery roots restricted to ComfyUI directories and validated wherever they are read; media-only, fail-closed image serving; Windows path escape in static routes closed; quote-safe HTML escaping with inline handlers moved to data attributes; the CivitAI key is never returned and only sent to HTTPS CivitAI hosts; auto-tag paths contained; side-effect routes are POST-only; backups include the WAL and restores are size-capped, integrity-checked and applied through SQLite's online backup API +- **🧡 Worker Threads**: a Settings slider (default half the cores) used by the image rescan, thumbnail generation and duplicate detection. Measured: a 1 GB sample scanned in 2.5 s instead of 18 s at 8 threads +- **πŸ” Background Jobs**: the image rescan and the duplicate scan run on the server, stream progress, survive a closed tab, minimize to a progress pill and re-attach after a reload +- **🎬 Video Prompts**: the scan reads the prompt ComfyUI's video savers embed in MP4/MOV/WebM containers (needs `ffprobe`) +- **πŸ–ΌοΈ Thumbnails**: written atomically and served under versioned URLs so a regenerated thumbnail never shows a stale cached copy; the per-run cap applies only to thumbnails that still need generating and the gallery says how many remain +- **🏷️ Tag Autocomplete**: the dashboard tag box suggests tags, narrowed to those that co-occur with the tags already typed +- **♾️ Infinite Scrolling**: optional in the gallery and the dashboard list (saved default plus a session toggle) +- **βš™οΈ Gallery Settings** now apply at load; grid columns, media info and the Rescan button work as labelled +- **πŸ› Correctness**: run counts increment when ComfyUI dequeues a prompt; negative prompts are never traced; per-thread SQLite connections and idempotent migrations (an old rebuild could delete every image row); `truncate_logs` no longer rolls the log over; the ComfyUI canvas loads only the canvas extension; Windows console banners are ASCII-safe; subfolders use `/` on every platform +- **πŸ§ͺ Tests & CI**: 1,460 Python and 142 JavaScript tests (from 510 and 26), coverage 96% with an 85% gate, CI on Linux, Windows and macOS for Python 3.10–3.12, flake8 clean +- **⚠️ API changes**: image paths in responses are relative to the ComfyUI tree (use `url`); `autotag/download`, `autotag/start` and the thumbnail progress stream are POST-only; `lora/status` returns `has_civitai_api_key` instead of the key; `/search` is paginated + +### v3.2.4 (Recently Used Sorting & Queue-Time Run Counts) + +- **πŸ•’ Recently Used Sort**: sort the dashboard by last use, most used or rating, server-side across pages +- **πŸ”’ Run Counts**: `run_count` and `last_used_at` are recorded through the ComfyUI queue hook, so cached nodes still count +- **πŸ”— Positive-Prompt Image Linking**: generated images link to the positive prompt traced through the queued graph + ### v3.2.1 (LoRA Manager Integration) - **πŸ”— LoRA Manager Integration**: Import LoRA metadata, trigger words, and CivitAI example images from [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manager) into your prompt database diff --git a/__init__.py b/__init__.py index 2e07092..e33aa42 100644 --- a/__init__.py +++ b/__init__.py @@ -5,7 +5,8 @@ gallery system. This module provides two main node types: - PromptManager: CLIP encoding node that outputs CONDITIONING -- PromptManagerText: Text-only node that outputs STRING with prepend/append functionality +- PromptManagerText: Text-only node that outputs STRING with prepend/append + functionality Both nodes share the same database backend for persistent prompt storage and include an automatic image gallery system that monitors ComfyUI output directories and links @@ -37,9 +38,9 @@ def get_version(): return "unknown" -from .prompt_manager import PromptManager -from .prompt_manager_text import PromptManagerText -from .prompt_search_list import PromptSearchList +from .prompt_manager import PromptManager # noqa: E402 +from .prompt_manager_text import PromptManagerText # noqa: E402 +from .prompt_search_list import PromptSearchList # noqa: E402 NODE_CLASS_MAPPINGS = { "PromptManager": PromptManager, @@ -53,8 +54,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "PromptSearchList": "Prompt Search List", } -# Define path to web directory for UI components -WEB_DIRECTORY = "web" +# Canvas extension only. ComfyUI loads every .js under this directory into the +# canvas page, so the admin/gallery bundles and vendored libraries under web/ +# are served by our own routes instead (see py/api). +WEB_DIRECTORY = "web/comfy" # Add API routes (same pattern as ComfyUI_Assets) try: @@ -82,6 +85,10 @@ except Exception as e: except Exception: pass +# Shared database, set up with the monitoring system below. Stays None when that +# fails so the queue hook can report the database as unavailable instead of crashing. +_global_db = None + # Start image monitoring globally at module import time # This ensures images are linked regardless of which node is used in the workflow try: @@ -111,13 +118,21 @@ except Exception as e: except Exception: print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}") -# Count prompt runs when workflows are queued (ComfyUI skips unchanged nodes, so + +def _db_factory(): + """The shared database for the queue hook; called lazily per request.""" + if _global_db is None: + raise RuntimeError("PromptManager database unavailable") + return _global_db + + +# Count prompt runs when ComfyUI dequeues workflows (it skips unchanged nodes, so # node execution alone misses re-runs). Separate from monitoring so either can fail. try: from .py.config import server_instance from .utils.usage_tracking import register_queue_hook - register_queue_hook(server_instance, lambda: _global_db) + register_queue_hook(server_instance, _db_factory) except Exception as e: try: from .utils.logging_config import get_logger @@ -136,6 +151,8 @@ __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] print() print(f"\033[94m[ComfyUI-PromptManager] Version:\033[0m {get_version()}") for node_key, display_name in NODE_DISPLAY_NAME_MAPPINGS.items(): - print(f"🫢 \033[94mLoaded:\033[0m {display_name}") + # ASCII only: a Windows console using a legacy code page cannot print emoji, + # and a UnicodeEncodeError here would abort loading the whole node. + print(f" \033[94mLoaded:\033[0m {display_name}") print(f"\033[94mTotal: {len(NODE_CLASS_MAPPINGS)} tools loaded\033[0m") print() diff --git a/database/models.py b/database/models.py index 5dc2cfe..84e7392 100644 --- a/database/models.py +++ b/database/models.py @@ -2,11 +2,16 @@ Database schema and models for KikoTextEncode prompt storage. """ +import contextlib import datetime +import shutil import sqlite3 import os import threading -from typing import Optional +import weakref +from contextlib import closing +from pathlib import Path +from typing import Set, Tuple # Import logging system try: @@ -25,15 +30,46 @@ ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00" def utc_now_iso() -> str: - """Current UTC time as ISO text with microseconds, e.g. 2026-10-07T18:39:06.758123+00:00.""" + """Current UTC time as ISO text with microseconds. + + Example: 2026-10-07T18:39:06.758123+00:00 + """ return datetime.datetime.now(datetime.timezone.utc).isoformat( timespec="microseconds" ) +def normalize_image_path(path: str) -> str: + """Canonical form of an image path for generated_images.file_path. + + Absolute, with '..' and '.' collapsed and the platform's case folding + applied, so two spellings of one file (or the same basename in two date + folders) compare the way the filesystem would. + """ + return os.path.normcase(os.path.normpath(os.path.abspath(path))) + + +# How long a connection waits for a locked database before giving up. Thirty +# seconds covers a scan or a thumbnail run writing while ComfyUI saves images, +# and the slow disks of CI runners. +BUSY_TIMEOUT_SECONDS = 30 +SQLITE_HEADER = b"SQLite format 3\x00" +REQUIRED_PROMPT_COLUMNS = ("id", "text", "created_at") + + class PromptModel: """Database model for prompt storage and schema management.""" + # Paths (normalised, see _schema_key) whose schema has been created and + # migrated in this process. Many PromptDatabase instances point at the + # same file; only the first one pays for the schema check. + _initialized_paths: Set[str] = set() + _init_lock = threading.Lock() + # Every live model, so a shutdown or a test teardown can close them all + # before the database file is deleted (Windows refuses otherwise). + _instances: "weakref.WeakSet[PromptModel]" = weakref.WeakSet() + _instances_lock = threading.Lock() + def __init__(self, db_path: str): """ Initialize the database model. @@ -44,8 +80,16 @@ class PromptModel: self.logger = get_logger("prompt_manager.database.models") self.logger.debug(f"Initializing database model with path: {db_path}") self.db_path = db_path - self._conn: Optional[sqlite3.Connection] = None + # One connection per thread (see get_connection); every open connection + # is also tracked here so restore_from_file can close them all. + self._local = threading.local() + self._connections: Set[sqlite3.Connection] = set() self._conn_lock = threading.Lock() + # Bumped by _close_all_connections so threads holding a closed + # connection reopen on their next get_connection(). + self._generation = 0 + with PromptModel._instances_lock: + PromptModel._instances.add(self) self._ensure_database_exists() def _ensure_database_exists(self) -> None: @@ -58,17 +102,59 @@ class PromptModel: Raises: Exception: If database creation fails """ + key = self._schema_key() + with PromptModel._init_lock: + if key in PromptModel._initialized_paths and os.path.exists(self.db_path): + return + try: + conn = sqlite3.connect(self.db_path, timeout=BUSY_TIMEOUT_SECONDS) + try: + conn.execute("PRAGMA journal_mode = WAL") + conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}") + conn.execute("PRAGMA foreign_keys = ON") + self._create_tables(conn) + self._create_indexes(conn) + conn.commit() + finally: + conn.close() + except Exception as e: + self.logger.error(f"Error creating database: {e}") + raise + PromptModel._initialized_paths.add(key) + + def _schema_key(self) -> str: + """Normalised identity of db_path for the once-per-process schema cache.""" + return os.path.normcase(os.path.realpath(self.db_path)) + + @classmethod + def reset_schema_cache(cls) -> None: + """Forget which databases were initialised (tests and restores).""" + with cls._init_lock: + cls._initialized_paths.clear() + + @contextlib.contextmanager + def _table_rebuild(self, conn: sqlite3.Connection): + """Run a copy-and-swap table rebuild with foreign keys disabled. + + DROP TABLE on a parent table with foreign_keys ON performs an implicit + DELETE that cascades into child rows (generated_images, prompt_tags), + so the rebuild runs with the pragma off. The pragma is a no-op inside + a transaction, hence the commits around it. Errors roll back and + propagate after the pragma is restored. + """ + conn.commit() + conn.execute("PRAGMA foreign_keys = OFF") + # Explicit BEGIN: in sqlite3's legacy transaction mode DDL autocommits, + # so without it a failed copy would leave the *_new table behind. + conn.execute("BEGIN") try: - with sqlite3.connect(self.db_path) as conn: - conn.execute("PRAGMA journal_mode = WAL") - conn.execute("PRAGMA busy_timeout = 5000") - conn.execute("PRAGMA foreign_keys = ON") - self._create_tables(conn) - self._create_indexes(conn) - conn.commit() - except Exception as e: - self.logger.error(f"Error creating database: {e}") + yield + conn.commit() + except BaseException: + conn.rollback() raise + finally: + conn.execute("PRAGMA foreign_keys = ON") def _create_tables(self, conn: sqlite3.Connection) -> None: """ @@ -112,8 +198,9 @@ class PromptModel: workflow_data TEXT, prompt_metadata TEXT, parameters TEXT, + file_path TEXT, FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE, - UNIQUE(prompt_id, filename) + UNIQUE(prompt_id, file_path) ) """) @@ -134,9 +221,6 @@ class PromptModel: ) """) - # Add unique constraint to existing databases (migration) - self._migrate_add_unique_constraint(conn) - # Check if we need to migrate from old schema with workflow_name self._migrate_workflow_name_removal(conn) @@ -146,9 +230,12 @@ class PromptModel: # Migrate JSON tags to normalized junction tables self._migrate_json_tags_to_junction(conn) - # Usage tracking columns (last, after any migration that rebuilds prompts) + # Usage tracking columns (after any migration that rebuilds prompts) self._migrate_add_usage_columns(conn) + # Image uniqueness by full path instead of basename (3.2.4) + self._migrate_image_uniqueness_by_path(conn) + def _create_indexes(self, conn: sqlite3.Connection) -> None: """ Create indexes for better query performance. @@ -163,43 +250,100 @@ class PromptModel: - Hash values for duplicate detection """ indexes = [ - "CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)", - "CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)", - "CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)", - "CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)", - "CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)", - "CREATE INDEX IF NOT EXISTS idx_prompts_last_used ON prompts(last_used_at)", - "CREATE INDEX IF NOT EXISTS idx_prompts_run_count ON prompts(run_count)", - "CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)", - "CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)", - "CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)", - "CREATE INDEX IF NOT EXISTS idx_prompt_tags_tag ON prompt_tags(tag_id)", - "CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name)", + ("idx_prompts_text", "prompts(text)"), + ("idx_prompts_category", "prompts(category)"), + ("idx_prompts_created_at", "prompts(created_at)"), + ("idx_prompts_hash", "prompts(hash)"), + ("idx_prompts_rating", "prompts(rating)"), + ("idx_prompts_last_used", "prompts(last_used_at)"), + ("idx_prompts_run_count", "prompts(run_count)"), + ("idx_prompt_images", "generated_images(prompt_id)"), + ("idx_image_path", "generated_images(image_path)"), + ("idx_image_file_path", "generated_images(file_path)"), + ("idx_generation_time", "generated_images(generation_time)"), + ("idx_prompt_tags_tag", "prompt_tags(tag_id)"), + ("idx_tags_name", "tags(name)"), ] - for index_sql in indexes: - conn.execute(index_sql) + for name, target in indexes: + conn.execute(f"CREATE INDEX IF NOT EXISTS {name} ON {target}") def get_connection(self) -> sqlite3.Connection: """ - Get the persistent database connection. + Get the calling thread's database connection. - Returns a thread-safe, reusable connection protected by a lock. - The connection is created on first call and reused thereafter. - Callers use it as a context manager; the lock is held for the - duration of the ``with`` block. + Each thread gets its own ``sqlite3.Connection`` (kept in + ``threading.local``), created on first use and reused by that thread + thereafter. Every connection is configured with WAL journaling, + ``row_factory = sqlite3.Row``, ``PRAGMA foreign_keys = ON`` and + ``PRAGMA busy_timeout`` of BUSY_TIMEOUT_SECONDS so concurrent writers + wait instead of failing. Callers use it as a context manager + (``with conn:``), which commits or rolls back that thread's transaction + and leaves the connection open. No lock is held while a connection is in use. Returns: - sqlite3.Connection: Configured database connection + sqlite3.Connection: The current thread's configured connection + """ + conn = getattr(self._local, "conn", None) + if conn is not None and self._local.generation == self._generation: + return conn + + # check_same_thread=False only so restore_from_file can close other + # threads' connections; each connection is still used by one thread. + conn = sqlite3.connect( + self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS + ) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode = WAL") + conn.execute("PRAGMA foreign_keys = ON") + conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}") + self._local.conn = conn + self._local.generation = self._generation + with self._conn_lock: + self._connections.add(conn) + return conn + + @classmethod + def close_all_instances(cls) -> None: + """Close every connection of every live model in this process.""" + with cls._instances_lock: + models = list(cls._instances) + for model in models: + model.close_all() + + def close_all(self) -> None: + """Close every thread's connection; they reopen lazily on next use. + + Needed before deleting the database file (Windows keeps open files + locked) and by tests that ran queries through executor threads. """ with self._conn_lock: - if self._conn is None: - self._conn = sqlite3.connect(self.db_path, check_same_thread=False) - self._conn.row_factory = sqlite3.Row - self._conn.execute("PRAGMA journal_mode = WAL") - self._conn.execute("PRAGMA foreign_keys = ON") - self._conn.execute("PRAGMA busy_timeout = 5000") - return self._conn + connections = list(self._connections) + self._connections.clear() + self._generation += 1 + self._local.conn = None + for conn in connections: + try: + conn.close() + except sqlite3.Error as e: + self.logger.warning(f"Error closing database connection: {e}") + + def close(self) -> None: + """Close the calling thread's connection, if it has one. + + Other threads keep their connections. Call this before deleting or + replacing the database file: Windows refuses to delete an open file. + """ + conn = getattr(self._local, "conn", None) + if conn is None: + return + self._local.conn = None + with self._conn_lock: + self._connections.discard(conn) + try: + conn.close() + except sqlite3.Error as e: + self.logger.warning(f"Error closing database connection: {e}") def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None: """ @@ -210,16 +354,20 @@ class PromptModel: This migration handles legacy schema updates by removing the deprecated workflow_name column while preserving all other data. + + Raises: + sqlite3.Error: If the rebuild fails (logged, nothing half-applied) """ try: - # Check if workflow_name column exists cursor = conn.execute("PRAGMA table_info(prompts)") columns = [column[1] for column in cursor.fetchall()] + if "workflow_name" not in columns: + return - if "workflow_name" in columns: - self.logger.info("Migrating database: removing workflow_name column") - - # Create new table without workflow_name + self.logger.info("Migrating database: removing workflow_name column") + with self._table_rebuild(conn): + # A previous attempt may have died between CREATE and RENAME + conn.execute("DROP TABLE IF EXISTS prompts_new") conn.execute(""" CREATE TABLE prompts_new ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -233,23 +381,20 @@ class PromptModel: hash TEXT UNIQUE ) """) - - # Copy data from old table to new table conn.execute(""" - INSERT INTO prompts_new (id, text, created_at, updated_at, category, tags, rating, notes, hash) - SELECT id, text, created_at, updated_at, category, tags, rating, notes, hash + INSERT INTO prompts_new (id, text, created_at, updated_at, + category, tags, rating, notes, hash) + SELECT id, text, created_at, updated_at, category, tags, rating, + notes, hash FROM prompts """) - - # Drop old table and rename new one conn.execute("DROP TABLE prompts") conn.execute("ALTER TABLE prompts_new RENAME TO prompts") + self.logger.info("Database migration completed") - self.logger.info("Database migration completed") - - except Exception as e: - self.logger.error(f"Migration error: {e}") - # If migration fails, the table creation will handle it + except sqlite3.Error as e: + self.logger.error(f"Migration error (workflow_name removal): {e}") + raise def _migrate_foreign_key_types(self, conn: sqlite3.Connection) -> None: """ @@ -260,18 +405,19 @@ class PromptModel: Converts prompt_id from TEXT to INTEGER type to match the prompts table's primary key type, ensuring referential integrity. + + Raises: + sqlite3.Error: If the rebuild fails (logged, nothing half-applied) """ try: - # Check if generated_images table exists and has TEXT prompt_id cursor = conn.execute("PRAGMA table_info(generated_images)") columns = {column[1]: column[2] for column in cursor.fetchall()} + if columns.get("prompt_id") != "TEXT": + return - if "prompt_id" in columns and columns["prompt_id"] == "TEXT": - self.logger.info( - "Migrating foreign key types: prompt_id TEXT -> INTEGER" - ) - - # Create new table with correct types + self.logger.info("Migrating foreign key types: prompt_id TEXT -> INTEGER") + with self._table_rebuild(conn): + conn.execute("DROP TABLE IF EXISTS generated_images_new") conn.execute(""" CREATE TABLE generated_images_new ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -289,130 +435,27 @@ class PromptModel: FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE ) """) - - # Copy data, converting prompt_id from TEXT to INTEGER conn.execute(""" - INSERT INTO generated_images_new - (id, prompt_id, image_path, filename, generation_time, file_size, + INSERT INTO generated_images_new + (id, prompt_id, image_path, filename, generation_time, file_size, width, height, format, workflow_data, prompt_metadata, parameters) - SELECT id, CAST(prompt_id AS INTEGER), image_path, filename, generation_time, - file_size, width, height, format, workflow_data, prompt_metadata, parameters + SELECT id, CAST(prompt_id AS INTEGER), image_path, filename, + generation_time, + file_size, width, height, format, workflow_data, + prompt_metadata, parameters FROM generated_images WHERE prompt_id != '' AND prompt_id IS NOT NULL AND CAST(prompt_id AS INTEGER) IN (SELECT id FROM prompts) """) - - # Drop old table and rename new one conn.execute("DROP TABLE generated_images") conn.execute( "ALTER TABLE generated_images_new RENAME TO generated_images" ) + self.logger.info("Foreign key migration completed") - self.logger.info("Foreign key migration completed") - - except Exception as e: - self.logger.error(f"Foreign key migration error: {e}") - # If migration fails, continue with existing schema - - def _migrate_add_unique_constraint(self, conn: sqlite3.Connection) -> None: - """ - Add UNIQUE constraint on (prompt_id, filename) to prevent duplicate image entries. - - This migration: - 1. Checks if the constraint already exists - 2. Removes duplicate entries (keeping the most recent) - 3. Recreates the table with the UNIQUE constraint - - Args: - conn: Active database connection - """ - try: - # Check if the unique constraint already exists by looking at table info - cursor = conn.execute("PRAGMA index_list(generated_images)") - indexes = cursor.fetchall() - - # Check if we have a unique index on prompt_id, filename - has_unique_constraint = False - for idx in indexes: - if idx[2] == 1: # unique flag - cursor = conn.execute(f"PRAGMA index_info({idx[1]})") - columns = [col[2] for col in cursor.fetchall()] - if "prompt_id" in columns and "filename" in columns: - has_unique_constraint = True - break - - if has_unique_constraint: - return # Already migrated - - self.logger.info( - "Migrating database: adding UNIQUE constraint on (prompt_id, filename)" - ) - - # First, remove duplicates keeping only the most recent (highest id) - conn.execute(""" - DELETE FROM generated_images - WHERE id NOT IN ( - SELECT MAX(id) FROM generated_images - GROUP BY prompt_id, filename - ) - """) - - duplicates_removed = conn.total_changes - if duplicates_removed > 0: - self.logger.info( - f"Removed {duplicates_removed} duplicate image entries" - ) - - # Create new table with UNIQUE constraint - conn.execute(""" - CREATE TABLE IF NOT EXISTS generated_images_new ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - prompt_id INTEGER NOT NULL, - image_path TEXT NOT NULL, - filename TEXT NOT NULL, - generation_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - file_size INTEGER, - width INTEGER, - height INTEGER, - format TEXT, - workflow_data TEXT, - prompt_metadata TEXT, - parameters TEXT, - FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE, - UNIQUE(prompt_id, filename) - ) - """) - - # Copy data - conn.execute(""" - INSERT INTO generated_images_new - (id, prompt_id, image_path, filename, generation_time, file_size, - width, height, format, workflow_data, prompt_metadata, parameters) - SELECT id, prompt_id, image_path, filename, generation_time, file_size, - width, height, format, workflow_data, prompt_metadata, parameters - FROM generated_images - """) - - # Drop old table and rename new one - conn.execute("DROP TABLE generated_images") - conn.execute("ALTER TABLE generated_images_new RENAME TO generated_images") - - # Recreate indexes - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)" - ) - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)" - ) - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)" - ) - - self.logger.info("UNIQUE constraint migration completed successfully") - - except Exception as e: - self.logger.error(f"UNIQUE constraint migration error: {e}") - # Continue with existing schema if migration fails + except sqlite3.Error as e: + self.logger.error(f"Migration error (foreign key types): {e}") + raise def _migrate_json_tags_to_junction(self, conn: sqlite3.Connection) -> None: """ @@ -420,6 +463,9 @@ class PromptModel: Runs once: skips if the tags table already has data. Uses json_each() to extract tag names from the JSON arrays stored in prompts.tags. + + Raises: + sqlite3.Error: If the migration fails (logged) """ try: cursor = conn.execute("SELECT COUNT(*) FROM tags") @@ -439,7 +485,8 @@ class PromptModel: conn.execute( "INSERT OR IGNORE INTO tags (name) " "SELECT DISTINCT je.value FROM prompts, json_each(prompts.tags) AS je " - "WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND prompts.tags != '[]'" + "WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND " + "prompts.tags != '[]'" ) # Populate junction table @@ -454,19 +501,22 @@ class PromptModel: tag_count = conn.execute("SELECT COUNT(*) FROM tags").fetchone()[0] link_count = conn.execute("SELECT COUNT(*) FROM prompt_tags").fetchone()[0] self.logger.info( - f"Tag migration complete: {tag_count} unique tags, {link_count} prompt-tag links" + f"Tag migration complete: {tag_count} unique tags, {link_count} " + "prompt-tag links" ) - except Exception as e: - self.logger.error(f"Tag junction migration error: {e}") + except sqlite3.Error as e: + self.logger.error(f"Migration error (tag junction): {e}") + raise def _migrate_add_usage_columns(self, conn: sqlite3.Connection) -> None: """ - Add last_used_at and run_count (3.2.4) and fill in any prompt without usage data. + Add last_used_at and run_count (3.2.4) and estimate usage for existing rows. - The backfill runs on every start but only touches rows whose last_used_at - is NULL, so it is a no-op once healed. That also repairs rows written by an - older version after a downgrade, or left unfilled by an interrupted upgrade. + The backfill runs only when the columns are added, i.e. on the upgrade + from a version without usage tracking. Afterwards a NULL last_used_at + means "never run": saving a prompt does not stamp it, only + record_prompt_use does, and "Recently Used" relies on that. The estimate comes from linked images: run_count is the image count (at least 1) and last_used_at the newest image time, falling back to created_at. @@ -475,14 +525,20 @@ class PromptModel: """ try: columns = {row[1] for row in conn.execute("PRAGMA table_info(prompts)")} + added = False if "last_used_at" not in columns: self.logger.info("Adding prompt usage column last_used_at") conn.execute("ALTER TABLE prompts ADD COLUMN last_used_at TIMESTAMP") + added = True if "run_count" not in columns: self.logger.info("Adding prompt usage column run_count") conn.execute( - "ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL DEFAULT 0" + "ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL " + "DEFAULT 0" ) + added = True + if not added: + return cursor = conn.execute(f""" UPDATE prompts SET @@ -501,18 +557,86 @@ class PromptModel: """) if cursor.rowcount: self.logger.info(f"Estimated usage for {cursor.rowcount} prompts") - except Exception as e: - self.logger.error(f"Usage tracking migration error: {e}") + except sqlite3.Error as e: + self.logger.error(f"Migration error (usage columns): {e}") + raise - def migrate_database(self) -> None: + def _migrate_image_uniqueness_by_path(self, conn: sqlite3.Connection) -> None: """ - Apply any pending database migrations. + Rebuild generated_images so an image is unique per (prompt_id, file_path). - This method serves as an entry point for future schema migrations. - Add new migration logic here as the database evolves. + Older schemas were unique per (prompt_id, filename), i.e. the basename, + so ComfyUI's same-named files in different date folders were silently + dropped. file_path holds normalize_image_path(image_path); rows whose + paths normalise to the same file collapse to the newest one. + + Raises: + sqlite3.Error: If the rebuild fails (logged, nothing half-applied) """ - # Future migrations can be added here - pass + try: + columns = { + row[1] for row in conn.execute("PRAGMA table_info(generated_images)") + } + if "file_path" in columns: + return + + self.logger.info("Migrating database: image uniqueness by full path") + with self._table_rebuild(conn): + rows = conn.execute( + "SELECT id, prompt_id, image_path FROM generated_images ORDER BY id" + ).fetchall() + keep = {} + for row_id, prompt_id, image_path in rows: + keep[(prompt_id, normalize_image_path(image_path))] = row_id + + conn.execute("DROP TABLE IF EXISTS generated_images_new") + conn.execute(""" + CREATE TABLE generated_images_new ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + prompt_id INTEGER NOT NULL, + image_path TEXT NOT NULL, + filename TEXT NOT NULL, + generation_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + file_size INTEGER, + width INTEGER, + height INTEGER, + format TEXT, + workflow_data TEXT, + prompt_metadata TEXT, + parameters TEXT, + file_path TEXT, + FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE + CASCADE, + UNIQUE(prompt_id, file_path) + ) + """) + conn.executemany( + """ + INSERT INTO generated_images_new + (id, prompt_id, image_path, filename, generation_time, file_size, + width, height, format, workflow_data, prompt_metadata, parameters, + file_path) + SELECT id, prompt_id, image_path, filename, generation_time, + file_size, + width, height, format, workflow_data, prompt_metadata, + parameters, + ? + FROM generated_images WHERE id = ? + """, + [(file_path, row_id) for (_, file_path), row_id in keep.items()], + ) + conn.execute("DROP TABLE generated_images") + conn.execute( + "ALTER TABLE generated_images_new RENAME TO generated_images" + ) + dropped = len(rows) - len(keep) + if dropped: + self.logger.info(f"Collapsed {dropped} duplicate image rows") + self.logger.info("Image path uniqueness migration completed") + + except sqlite3.Error as e: + self.logger.error(f"Migration error (image path uniqueness): {e}") + raise def vacuum_database(self) -> None: """ @@ -522,7 +646,7 @@ class PromptModel: improving query performance and reducing file size. """ try: - with sqlite3.connect(self.db_path) as conn: + with closing(sqlite3.connect(self.db_path)) as conn: conn.execute("VACUUM") conn.commit() except Exception as e: @@ -541,12 +665,14 @@ class PromptModel: total_prompts = cursor.fetchone()["total_prompts"] cursor = conn.execute( - "SELECT COUNT(DISTINCT category) as unique_categories FROM prompts WHERE category IS NOT NULL" + "SELECT COUNT(DISTINCT category) as unique_categories FROM " + "prompts WHERE category IS NOT NULL" ) unique_categories = cursor.fetchone()["unique_categories"] cursor = conn.execute( - "SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS NOT NULL" + "SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS " + "NOT NULL" ) avg_rating = cursor.fetchone()["avg_rating"] @@ -568,19 +694,156 @@ class PromptModel: def backup_database(self, backup_path: str) -> bool: """ - Create a backup of the database. + Write a consistent copy of the database to backup_path. + + Uses the SQLite online backup API through this thread's connection, + so rows that are committed but still sit in the -wal file are + included and other threads may keep reading and writing meanwhile. Args: - backup_path: Path where the backup should be saved + backup_path: Path where the backup should be saved (overwritten) Returns: bool: True if backup was successful, False otherwise """ try: - import shutil - - shutil.copy2(self.db_path, backup_path) + if os.path.exists(backup_path) and os.path.samefile( + backup_path, self.db_path + ): + self.logger.error("Refusing to back up the database onto itself") + return False + if os.path.exists(backup_path): + os.remove(backup_path) + source = self.get_connection() + destination = sqlite3.connect(backup_path) + try: + source.backup(destination) + finally: + destination.close() return True - except Exception as e: + except (OSError, sqlite3.Error) as e: self.logger.error(f"Error creating database backup: {e}") return False + + @staticmethod + def verify_database_file(path: str) -> Tuple[bool, str]: + """ + Check that a file is a healthy prompts database before restoring it. + + Opens the file read-only, runs PRAGMA integrity_check and confirms the + prompts table has the columns every version of this extension needs. + + Args: + path: File to inspect + + Returns: + (ok, reason): reason is "ok" when the file is acceptable, otherwise + a human-readable explanation. Never raises for bad input. + """ + if not os.path.isfile(path): + return False, "File not found" + try: + with open(path, "rb") as fh: + header = fh.read(len(SQLITE_HEADER)) + except OSError as e: + return False, f"Cannot read file: {e}" + if header != SQLITE_HEADER: + return False, "Not a SQLite database file" + + uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro" + try: + conn = sqlite3.connect(uri, uri=True) + except sqlite3.Error as e: + return False, f"Invalid SQLite database: {e}" + try: + row = conn.execute("PRAGMA integrity_check").fetchone() + verdict = row[0] if row else "no result" + if verdict != "ok": + return False, f"Integrity check failed: {verdict}" + table = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'" + ).fetchone() + if not table: + return False, "Database does not contain a 'prompts' table" + columns = {r[1] for r in conn.execute("PRAGMA table_info(prompts)")} + missing = [c for c in REQUIRED_PROMPT_COLUMNS if c not in columns] + if missing: + return ( + False, + "prompts table is missing required columns: " + ", ".join(missing), + ) + return True, "ok" + except sqlite3.Error as e: + return False, f"Invalid SQLite database: {e}" + finally: + conn.close() + + def _replace_contents_from(self, src_path: str) -> None: + """Copy every page of *src_path* into the live database. + + Goes through SQLite's online backup API into this thread's normal + connection, so SQLite takes the locks itself and the other + PromptModel instances that share the file (nodes, image monitor, + queue hook, API) simply see the new pages on their next statement. + Nothing is swapped underneath open handles and no -wal/-shm file has + to be deleted, which also keeps Windows happy. When SQLite refuses + (page sizes differ while the target is in WAL mode) fall back to a + file copy with every connection of this instance closed first. + """ + source = sqlite3.connect( + Path(src_path).resolve().as_uri() + "?mode=ro", uri=True + ) + try: + try: + source.backup(self.get_connection()) + return + except sqlite3.OperationalError as e: + self.logger.warning( + f"Online restore failed ({e}); falling back to a file copy" + ) + finally: + source.close() + self.close_all() + shutil.copy2(src_path, self.db_path) + for suffix in ("-wal", "-shm"): + stale = self.db_path + suffix + if os.path.exists(stale): + os.remove(stale) + + def restore_from_file(self, src_path: str) -> str: + """ + Replace the live database with the file at src_path. + + The current database is first backed up next to itself as + ``.backup_`` (WAL content included), the + pages of ``src_path`` are copied into the live database through + SQLite's backup API (visible to every open connection) and the schema + migrations run on the restored content. + + Callers should run verify_database_file(src_path) first. + + Args: + src_path: SQLite file to restore + + Returns: + Path of the backup taken of the previous database, or "" when + there was no database to back up. + + Raises: + RuntimeError: If the previous database could not be backed up + OSError: If the restored file could not be copied into place + """ + backup_path = "" + if os.path.exists(self.db_path): + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + backup_path = f"{self.db_path}.backup_{timestamp}" + if not self.backup_database(backup_path): + raise RuntimeError("Could not back up the current database") + self.logger.info(f"Current database backed up to: {backup_path}") + + self._replace_contents_from(src_path) + with PromptModel._init_lock: + PromptModel._initialized_paths.discard(self._schema_key()) + self._ensure_database_exists() + self.logger.info(f"Database restored from {src_path}") + return backup_path diff --git a/database/operations.py b/database/operations.py index f3c0765..a1b8bcb 100644 --- a/database/operations.py +++ b/database/operations.py @@ -6,9 +6,9 @@ import sqlite3 import json import datetime import os -from typing import Optional, List, Dict, Any, Union +from typing import Optional, List, Dict, Any -from .models import PromptModel, utc_now_iso +from .models import PromptModel, normalize_image_path, utc_now_iso # Import logging system try: @@ -30,11 +30,15 @@ TAG_SUBQUERY = ( # Server-side sort orders, keyed by the value the UI sends. Only these strings ever # reach ORDER BY; unknown keys fall back to DEFAULT_SORT. +# last_used_at is NULL until the prompt has run, so "IS NULL" first keeps +# never-run prompts after every run one (SQLite sorts NULL first in DESC). SORT_ORDERS = { - "last_used_desc": "last_used_at DESC, id DESC", + "last_used_desc": "last_used_at IS NULL, last_used_at DESC, id DESC", "created_desc": "created_at DESC, id DESC", "created_asc": "created_at ASC, id ASC", - "run_count_desc": "run_count DESC, last_used_at DESC, id DESC", + "run_count_desc": ( + "run_count DESC, last_used_at IS NULL, last_used_at DESC, id DESC" + ), "rating_desc": "rating IS NULL, rating DESC, id DESC", "rating_asc": "rating IS NULL, rating ASC, id DESC", "text_asc": "text COLLATE NOCASE ASC, id ASC", @@ -48,6 +52,18 @@ def order_by_clause(sort: Optional[str]) -> str: return SORT_ORDERS.get(sort or DEFAULT_SORT, SORT_ORDERS[DEFAULT_SORT]) +LIKE_ESCAPE = "\\" + + +def escape_like(term: str) -> str: + """Escape LIKE wildcards in user text; use with ``LIKE ? ESCAPE '\\'``.""" + return ( + term.replace(LIKE_ESCAPE, LIKE_ESCAPE + LIKE_ESCAPE) + .replace("%", LIKE_ESCAPE + "%") + .replace("_", LIKE_ESCAPE + "_") + ) + + def _resolve_db_path(db_path: Optional[str] = None) -> str: """Resolve the database path from config, falling back to defaults. @@ -73,7 +89,7 @@ def _resolve_db_path(db_path: Optional[str] = None) -> str: db_path = "prompts.db" if not os.path.isabs(db_path): - db_path = os.path.join(extension_root, db_path) + db_path = os.path.normpath(os.path.join(extension_root, db_path)) # Ensure parent directory exists for custom paths parent = os.path.dirname(db_path) @@ -83,6 +99,51 @@ def _resolve_db_path(db_path: Optional[str] = None) -> str: return db_path +# Bounds for the file_info a link request may attach to an image. Anything +# outside them (or of the wrong type) is stored as NULL instead of being +# trusted into an INTEGER or TEXT column and echoed back to the gallery. +MAX_IMAGE_FILE_SIZE = 10**12 +MAX_IMAGE_DIMENSION = 65535 +IMAGE_FORMATS = frozenset( + {"PNG", "JPEG", "JPG", "WEBP", "GIF", "BMP", "TIFF", "MP4", "WEBM", "MOV"} +) + + +def _bounded_int(value: Any, upper: int) -> Optional[int]: + """``value`` as an int within ``[0, upper]``; None for anything else.""" + if isinstance(value, bool): + return None + try: + number = int(value) + except (TypeError, ValueError, OverflowError): + return None + return number if 0 <= number <= upper else None + + +def _known_image_format(value: Any) -> Optional[str]: + """Upper-cased *value* when it names a supported media format, else None.""" + if not isinstance(value, str): + return None + name = value.strip().upper() + return name if name in IMAGE_FORMATS else None + + +def _coerce_file_info(metadata: Any) -> tuple: + """``(size, width, height, format)`` from request metadata, each validated.""" + file_info = metadata.get("file_info") if isinstance(metadata, dict) else None + if not isinstance(file_info, dict): + return None, None, None, None + dimensions = file_info.get("dimensions") + if not isinstance(dimensions, (list, tuple)) or len(dimensions) < 2: + dimensions = (None, None) + return ( + _bounded_int(file_info.get("size"), MAX_IMAGE_FILE_SIZE), + _bounded_int(dimensions[0], MAX_IMAGE_DIMENSION), + _bounded_int(dimensions[1], MAX_IMAGE_DIMENSION), + _known_image_format(file_info.get("format")), + ) + + class PromptDatabase: """Database operations class for managing prompts.""" @@ -100,6 +161,19 @@ class PromptDatabase: self.model = PromptModel(db_path) self.logger.debug("Database operations initialized successfully") + def close(self) -> None: + """Close the calling thread's database connection (see PromptModel.close).""" + self.model.close() + + @staticmethod + def close_all_instances() -> None: + """Close the connections of every PromptDatabase/PromptModel in the process.""" + PromptModel.close_all_instances() + + def close_all(self) -> None: + """Close every thread's connection (see PromptModel.close_all).""" + self.model.close_all() + def save_prompt( self, text: str, @@ -134,34 +208,46 @@ class PromptDatabase: raise ValueError("Rating must be between 1 and 5") self.logger.debug( - f"Saving prompt: text_length={len(text)}, category={category}, tags={tags}, rating={rating}" + f"Saving prompt: text_length={len(text)}, category={category}, " + f"tags={tags}, rating={rating}" ) - with self.model.get_connection() as conn: - cursor = conn.execute( - """ - INSERT INTO prompts ( - text, category, tags, rating, notes, hash, created_at, updated_at, - last_used_at - ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?, ?) - """, - ( - text.strip(), - category, - rating, - notes, - prompt_hash, - datetime.datetime.now(datetime.timezone.utc).isoformat(), - datetime.datetime.now(datetime.timezone.utc).isoformat(), - utc_now_iso(), - ), + try: + with self.model.get_connection() as conn: + cursor = conn.execute( + """ + INSERT INTO prompts ( + text, category, tags, rating, notes, hash, created_at, + updated_at + ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?) + """, + ( + text.strip(), + category, + rating, + notes, + prompt_hash, + datetime.datetime.now(datetime.timezone.utc).isoformat(), + datetime.datetime.now(datetime.timezone.utc).isoformat(), + ), + ) + prompt_id = cursor.lastrowid + if tags: + self._sync_prompt_tags(conn, prompt_id, tags) + conn.commit() + self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}") + return prompt_id + except sqlite3.IntegrityError: + # Another caller inserted the same hash between the caller's + # duplicate check and this INSERT (the with-block rolled back). + existing = self.get_prompt_by_hash(prompt_hash) if prompt_hash else None + if existing is None: + raise + self.logger.debug( + f"Prompt with hash {prompt_hash} was saved concurrently; " + f"returning existing ID {existing['id']}" ) - prompt_id = cursor.lastrowid - if tags: - self._sync_prompt_tags(conn, prompt_id, tags) - conn.commit() - self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}") - return prompt_id + return existing["id"] def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool: """ @@ -258,8 +344,8 @@ class PromptDatabase: params = [] if text: - query_parts.append("AND text LIKE ?") - params.append(f"%{text}%") + query_parts.append("AND text LIKE ? ESCAPE '\\'") + params.append(f"%{escape_like(text)}%") if category: query_parts.append("AND category = ?") @@ -271,9 +357,10 @@ class PromptDatabase: query_parts.append( "AND prompts.id IN (" " SELECT pt.prompt_id FROM prompt_tags pt" - " JOIN tags t ON pt.tag_id = t.id WHERE t.name LIKE ?)" + " JOIN tags t ON pt.tag_id = t.id" + " WHERE t.name LIKE ? ESCAPE '\\')" ) - params.append(f"%{tag}%") + params.append(f"%{escape_like(tag)}%") else: query_parts.append( "AND prompts.id IN (" @@ -392,7 +479,8 @@ class PromptDatabase: """ with self.model.get_connection() as conn: cursor = conn.execute( - f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE category = ? ORDER BY created_at DESC LIMIT ?", + f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE category = ? " + "ORDER BY created_at DESC LIMIT ?", (category, limit), ) rows = cursor.fetchall() @@ -468,7 +556,7 @@ class PromptDatabase: params.append(datetime.datetime.now(datetime.timezone.utc).isoformat()) params.append(prompt_id) query = f"UPDATE prompts SET {', '.join(updates)} WHERE id = ?" - cursor = conn.execute(query, params) + conn.execute(query, params) if tags is not None: self._sync_prompt_tags(conn, prompt_id, tags) conn.execute( @@ -542,7 +630,8 @@ class PromptDatabase: """ with self.model.get_connection() as conn: cursor = conn.execute( - "SELECT DISTINCT TRIM(category) as category FROM prompts WHERE category IS NOT NULL AND TRIM(category) != '' ORDER BY category" + "SELECT DISTINCT TRIM(category) as category FROM prompts WHERE " + "category IS NOT NULL AND TRIM(category) != '' ORDER BY category" ) return [row["category"] for row in cursor.fetchall()] @@ -583,7 +672,8 @@ class PromptDatabase: rel = os.path.relpath(parent, root) if not rel.startswith(".."): if rel != ".": - folders.add(rel) + # Clients split and filter on "/" whatever the OS + folders.add(rel.replace(os.sep, "/")) made_relative = True break except ValueError: @@ -593,7 +683,7 @@ class PromptDatabase: # Preserve full relative-style path instead of collapsing # to just the basename, which loses hierarchy and creates # ambiguity (e.g. foo/bar and baz/bar both become "bar"). - folders.add(parent) + folders.add(parent.replace(os.sep, "/")) if include_ancestors: ancestors = set() @@ -643,8 +733,8 @@ class PromptDatabase: search_clause = "" params: list = [] if search: - search_clause = "HAVING t.name LIKE ?" - params.append(f"%{search}%") + search_clause = "HAVING t.name LIKE ? ESCAPE '\\'" + params.append(f"%{escape_like(search)}%") sort_map = { "alpha_desc": "tag COLLATE NOCASE DESC", @@ -685,6 +775,53 @@ class PromptDatabase: "has_more": (offset + limit) < total, } + def suggest_tags( + self, prefix: str, with_tags: List[str], limit: int = 15 + ) -> List[Dict[str, Any]]: + """Tag autocomplete narrowed by co-occurrence. + + Returns tags whose name starts with ``prefix`` (case-insensitive), + counted over the prompts that carry every tag in ``with_tags``; the + context tags themselves are left out. Ordered by count, then name. + """ + seen = set() + context = [] + for raw in with_tags or []: + name = str(raw).strip() + if name and name.lower() not in seen: + seen.add(name.lower()) + context.append(name) + + where = ["t.name LIKE ? ESCAPE '\\'"] + params: list = [f"{escape_like(prefix or '')}%"] + if context: + placeholders = ",".join("?" * len(context)) + where.append( + "pt.prompt_id IN (" + " SELECT pt2.prompt_id FROM prompt_tags pt2" + " JOIN tags t2 ON t2.id = pt2.tag_id" + f" WHERE t2.name COLLATE NOCASE IN ({placeholders})" + " GROUP BY pt2.prompt_id" + " HAVING COUNT(DISTINCT lower(t2.name)) = ?)" + ) + params.extend(context) + params.append(len(context)) + where.append(f"t.name COLLATE NOCASE NOT IN ({placeholders})") + params.extend(context) + params.append(max(1, int(limit))) + + sql = ( + "SELECT t.name AS name, COUNT(*) AS count" + " FROM prompt_tags pt JOIN tags t ON t.id = pt.tag_id" + f" WHERE {' AND '.join(where)}" + " GROUP BY t.id" + " ORDER BY count DESC, t.name COLLATE NOCASE ASC" + " LIMIT ?" + ) + with self.model.get_connection() as conn: + rows = conn.execute(sql, params).fetchall() + return [{"name": row["name"], "count": row["count"]} for row in rows] + def get_prompts_by_tags( self, tags: List[str], mode: str = "and", limit: int = 20, offset: int = 0 ) -> Dict[str, Any]: @@ -917,7 +1054,8 @@ class PromptDatabase: conn.commit() self.logger.info( - f"Merged {tags_merged} tags into '{target_tag}', affected {affected} prompts" + f"Merged {tags_merged} tags into '{target_tag}', affected {affected} " + "prompts" ) return { "success": True, @@ -936,7 +1074,8 @@ class PromptDatabase: with self.model.get_connection() as conn: cursor = conn.execute( "SELECT COUNT(*) as total FROM prompts " - "WHERE NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = prompts.id)" + "WHERE NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = " + "prompts.id)" ) return cursor.fetchone()["total"] @@ -1049,7 +1188,8 @@ class PromptDatabase: tag_id = tag_map.get(tag_name) if tag_id: conn.execute( - "INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)", + "INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) " + "VALUES (?, ?)", (prompt_id, tag_id), ) @@ -1138,79 +1278,6 @@ class PromptDatabase: self.logger.error(f"Error exporting prompts: {e}") return False - def find_duplicates(self) -> List[Dict[str, Any]]: - """ - Find duplicate prompts based on text content without removing them. - - Returns: - List of duplicate groups, each containing: - - text: The duplicate text content - - prompts: List of prompt records with same text - """ - self.logger.info("Scanning for duplicate prompts") - try: - with self.model.get_connection() as conn: - # Find duplicates by text content (case-insensitive) - # Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility - # We'll sort the IDs manually after fetching - cursor = conn.execute(""" - SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count, - GROUP_CONCAT(id) as ids, - GROUP_CONCAT(created_at) as created_dates - FROM prompts - GROUP BY LOWER(TRIM(text)) - HAVING COUNT(*) > 1 - """) - - duplicate_groups = cursor.fetchall() - self.logger.debug( - f"Found {len(duplicate_groups)} groups of duplicate prompts" - ) - - result = [] - - for group in duplicate_groups: - ids = group["ids"].split(",") - created_dates = group["created_dates"].split(",") - - # Sort IDs by created_at date - id_date_pairs = list(zip(ids, created_dates)) - id_date_pairs.sort(key=lambda x: x[1]) # Sort by date - ids = [pair[0] for pair in id_date_pairs] - - # Get full details for all prompts in this duplicate group - prompts = [] - for prompt_id in ids: - cursor = conn.execute( - "SELECT id, text, category, rating, created_at, updated_at " - "FROM prompts WHERE id = ?", - (int(prompt_id),), - ) - prompt_data = cursor.fetchone() - if prompt_data: - prompt_dict = dict(prompt_data) - prompt_dict["tags"] = self._get_prompt_tags( - conn, int(prompt_id) - ) - prompts.append(prompt_dict) - - if prompts: - result.append( - { - "text": prompts[0][ - "text" - ], # Use the actual text (not normalized) - "prompts": prompts, - } - ) - - self.logger.info(f"Found {len(result)} groups with duplicates") - return result - - except Exception as e: - self.logger.error(f"Error finding duplicates: {e}") - return [] - def cleanup_duplicates(self) -> int: """ Remove duplicate prompts based on text content, preserving all image links. @@ -1255,7 +1322,8 @@ class PromptDatabase: duplicate_ids = sorted_ids[1:] self.logger.debug( - f"Merging duplicates: keeping {primary_id}, removing {duplicate_ids}" + f"Merging duplicates: keeping {primary_id}, removing " + f"{duplicate_ids}" ) # Get primary prompt details @@ -1295,7 +1363,8 @@ class PromptDatabase: if total_removed > 0: self.logger.info( - f"Removed {total_removed} duplicate prompts, transferred {total_images_transferred} images" + f"Removed {total_removed} duplicate prompts, transferred " + f"{total_images_transferred} images" ) else: self.logger.info("No duplicate prompts found") @@ -1348,48 +1417,46 @@ class PromptDatabase: ) if not cursor.fetchone(): self.logger.warning( - f"Prompt ID {prompt_id_int} not found in database, skipping image linking" + f"Prompt ID {prompt_id_int} not found in database, skipping " + "image linking" ) return 0 # Proceed with linking + if not isinstance(metadata, dict): + metadata = {} filename = os.path.basename(image_path) - file_info = metadata.get("file_info", {}) if metadata else {} + file_path = normalize_image_path(image_path) + file_size, width, height, image_format = _coerce_file_info(metadata) - # Use INSERT OR IGNORE to skip duplicates (same prompt_id + filename) + # INSERT OR IGNORE skips a file already linked to this prompt + # (same prompt_id + normalised full path) cursor = conn.execute( """ INSERT OR IGNORE INTO generated_images - (prompt_id, image_path, filename, file_size, width, height, format, - workflow_data, prompt_metadata, parameters) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + (prompt_id, image_path, filename, file_path, file_size, width, + height, format, workflow_data, prompt_metadata, parameters) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( prompt_id_int, image_path, filename, - file_info.get("size"), - ( - file_info.get("dimensions", [None, None])[0] - if file_info.get("dimensions") - else None - ), - ( - file_info.get("dimensions", [None, None])[1] - if file_info.get("dimensions") - else None - ), - file_info.get("format"), - json.dumps(metadata.get("workflow", {}) if metadata else {}), - json.dumps(metadata.get("prompt", {}) if metadata else {}), - json.dumps(metadata.get("parameters", {}) if metadata else {}), + file_path, + file_size, + width, + height, + image_format, + json.dumps(metadata.get("workflow", {})), + json.dumps(metadata.get("prompt", {})), + json.dumps(metadata.get("parameters", {})), ), ) conn.commit() - if cursor.lastrowid == 0: + if cursor.rowcount == 0: self.logger.debug( - f"Image {filename} already linked to prompt {prompt_id_int}" + f"Image {file_path} already linked to prompt {prompt_id_int}" ) return 0 @@ -1420,12 +1487,15 @@ class PromptDatabase: ) return [self._image_row_to_dict(row) for row in cursor.fetchall()] - def get_recent_images(self, limit: int = 50) -> List[Dict[str, Any]]: + def get_recent_images( + self, limit: int = 50, offset: int = 0 + ) -> List[Dict[str, Any]]: """ - Get recently generated images across all prompts. + Get recently generated images across all prompts, newest first. Args: limit: Maximum number of images to return + offset: Number of newest images to skip (paging) Returns: List of image records with prompt text @@ -1437,9 +1507,9 @@ class PromptDatabase: FROM generated_images gi LEFT JOIN prompts p ON gi.prompt_id = p.id ORDER BY gi.generation_time DESC - LIMIT ? + LIMIT ? OFFSET ? """, - (limit,), + (int(limit), max(0, int(offset))), ) return [self._image_row_to_dict(row) for row in cursor.fetchall()] @@ -1461,7 +1531,8 @@ class PromptDatabase: sql = ( "SELECT gi.*, p.text as prompt_text, " "(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt " - "JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = p.id) AS _prompt_tags_list " + "JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = p.id) AS " + "_prompt_tags_list " "FROM generated_images gi " "INNER JOIN prompts p ON gi.prompt_id = p.id " "WHERE gi.image_path IS NOT NULL AND gi.image_path != '' " @@ -1485,12 +1556,16 @@ class PromptDatabase: result.append(data) return result - def search_images_by_prompt(self, search_term: str) -> List[Dict[str, Any]]: + def search_images_by_prompt( + self, search_term: str, limit: Optional[int] = None, offset: int = 0 + ) -> List[Dict[str, Any]]: """ - Search images by prompt text. + Search images by prompt text, newest first. Args: search_term: Text to search for in prompt content + limit: Maximum number of images to return (None for no limit) + offset: Number of newest matches to skip (paging) Returns: List of image records with prompt text @@ -1501,10 +1576,15 @@ class PromptDatabase: SELECT gi.*, p.text as prompt_text FROM generated_images gi JOIN prompts p ON gi.prompt_id = p.id - WHERE p.text LIKE ? + WHERE p.text LIKE ? ESCAPE '\\' ORDER BY gi.generation_time DESC + LIMIT ? OFFSET ? """, - (f"%{search_term}%",), + ( + f"%{escape_like(search_term)}%", + -1 if limit is None else int(limit), + max(0, int(offset)), + ), ) return [self._image_row_to_dict(row) for row in cursor.fetchall()] @@ -1675,7 +1755,8 @@ class PromptDatabase: if cursor.rowcount > 0: self.logger.debug( - f"Transferred {cursor.rowcount} images from prompt {dup_id} to {primary_id}" + f"Transferred {cursor.rowcount} images from prompt {dup_id} " + f"to {primary_id}" ) return transferred_count @@ -1853,7 +1934,8 @@ class PromptDatabase: tag_id = tag_map.get(tag_name) if tag_id: cursor = conn.execute( - "INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)", + "INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) " + "VALUES (?, ?)", (pid, tag_id), ) if cursor.rowcount > 0: @@ -1897,8 +1979,13 @@ class PromptDatabase: gi.workflow_data, gi.prompt_metadata, gi.generation_time FROM generated_images gi JOIN prompts p ON gi.prompt_id = p.id - WHERE gi.image_path = ? OR gi.image_path LIKE ?""", - (image_path, f"%{os.path.basename(image_path)}"), + WHERE gi.image_path = ? OR gi.file_path = ? + OR gi.image_path LIKE ? ESCAPE '\\'""", + ( + image_path, + normalize_image_path(image_path), + f"%{escape_like(os.path.basename(image_path))}", + ), ) row = cursor.fetchone() if not row: @@ -1928,8 +2015,9 @@ class PromptDatabase: """Return the prompt_id linked to an image, or None.""" with self.model.get_connection() as conn: cursor = conn.execute( - "SELECT prompt_id FROM generated_images WHERE image_path = ?", - (image_path,), + "SELECT prompt_id FROM generated_images" + " WHERE image_path = ? OR file_path = ?", + (image_path, normalize_image_path(image_path)), ) row = cursor.fetchone() return row[0] if row else None @@ -1984,7 +2072,8 @@ class PromptDatabase: """) for ref in cursor.fetchall(): issues.append( - f"prompt_tags entry references non-existent prompt {ref['prompt_id']}" + "prompt_tags entry references non-existent prompt " + f"{ref['prompt_id']}" ) # Check for orphaned image entries @@ -1995,7 +2084,8 @@ class PromptDatabase: """) for ref in cursor.fetchall(): issues.append( - f"Image {ref['id']} references non-existent prompt {ref['prompt_id']}" + f"Image {ref['id']} references non-existent prompt " + f"{ref['prompt_id']}" ) return issues diff --git a/prompt_manager.py b/prompt_manager.py index 9637c7c..90904c0 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -51,7 +51,9 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): { "multiline": True, "dynamicPrompts": True, - "tooltip": "The text prompt to be encoded and saved to database.", + "tooltip": ( + "The text prompt to be encoded and saved to database." + ), }, ), "clip": ( @@ -64,33 +66,48 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): IO.STRING, { "default": "", - "tooltip": "Optional category for organizing prompts (e.g., 'landscapes', 'portraits')", + "tooltip": ( + "Optional category for organizing prompts " + "(e.g., 'landscapes', 'portraits')" + ), }, ), "tags": ( IO.STRING, { "default": "", - "tooltip": "Comma-separated tags for the prompt (e.g., 'anime, detailed, sunset')", + "tooltip": ( + "Comma-separated tags for the prompt " + "(e.g., 'anime, detailed, sunset')" + ), }, ), "search_text": ( IO.STRING, { "default": "", - "tooltip": 'Search past prompts by text; add tag:name (or tag:"two words") to filter by tag', + "tooltip": ( + "Search past prompts by text; add tag:name " + '(or tag:"two words") to filter by tag' + ), }, ), "prepend_text": ( IO.STRING, { - "tooltip": "Text to prepend to the main prompt (connected STRING nodes will be added before the main text)" + "tooltip": ( + "Text to prepend to the main prompt (connected STRING " + "nodes will be added before the main text)" + ) }, ), "append_text": ( IO.STRING, { - "tooltip": "Text to append to the main prompt (connected STRING nodes will be added after the main text)" + "tooltip": ( + "Text to append to the main prompt (connected STRING " + "nodes will be added after the main text)" + ) }, ), }, @@ -104,15 +121,18 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): RETURN_TYPES = (IO.CONDITIONING, IO.STRING) OUTPUT_TOOLTIPS = ( - "A conditioning containing the embedded text used to guide the diffusion model.", - "The final combined text string (with prepend/append applied) that was encoded.", + "A conditioning containing the embedded text used to guide the diffusion " + "model.", + "The final combined text string (with prepend/append applied) that was " + "encoded.", ) FUNCTION = "encode_prompt" OUTPUT_NODE = True CATEGORY = "🫢 ComfyAssets/🧠 Prompts" DESCRIPTION = ( - "Encodes a text prompt using a CLIP model into an embedding that can be used to guide " - "the diffusion model towards generating specific images. Additionally saves all prompts " + "Encodes a text prompt using a CLIP model into an embedding that can be used " + "to guide the diffusion model towards generating specific images. " + "Additionally saves all prompts " "to a local SQLite database with optional metadata for search and retrieval." ) @@ -141,7 +161,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): append_text: Text to append to the main prompt Returns: - Tuple containing the conditioning for the diffusion model and the final text string + Tuple containing the conditioning for the diffusion model and the final + text string Raises: RuntimeError: If clip input is invalid @@ -162,7 +183,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): # Use the combined text for encoding encoding_text = final_text - # For database storage, save the original main text with metadata about prepend/append + # For database storage, save the original main text with metadata + # about prepend/append storage_text = text # Validate CLIP model @@ -251,10 +273,11 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): """ ComfyUI method to determine if node needs re-execution. - Returns a hash of input values that affect the conditioning output. - This enables proper branch execution - only re-execute when inputs change. + Returns a hash of input values that affect the conditioning output or + the saved prompt's metadata. This enables proper branch execution - + only re-execute when inputs change. """ import hashlib - combined = f"{text}|{prepend_text}|{append_text}" + combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}" return hashlib.sha256(combined.encode()).hexdigest() diff --git a/prompt_manager_base.py b/prompt_manager_base.py index ee190d5..2c8550f 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -6,9 +6,7 @@ search, gallery system management, and cleanup β€” extracted from the duplicate code in prompt_manager.py and prompt_manager_text.py. """ -import hashlib import os -import webbrowser from typing import Any, Dict, List, Optional try: @@ -225,16 +223,8 @@ class PromptManagerBase: return modified def _generate_hash(self, text: str) -> str: - """Generate SHA256 hash for the prompt text. - - Args: - text: The prompt text to hash - - Returns: - Hexadecimal string representation of the SHA256 hash - """ - normalized_text = text.strip().lower() - return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest() + """The prompt's deduplication hash (see utils.hashing.generate_prompt_hash).""" + return generate_prompt_hash(text) def _parse_tags(self, tags_string: str) -> Optional[list]: """Parse comma-separated tags string into a list. @@ -285,32 +275,6 @@ class PromptManagerBase: self.logger.error(f"Error searching prompts: {e}") return [] - def _open_web_interface(self): - """Open the web interface in the default browser.""" - try: - current_dir = os.path.dirname(os.path.abspath(__file__)) - web_dir = os.path.join(current_dir, "web_interface") - - if os.path.exists(web_dir): - index_path = os.path.join(web_dir, "index.html") - if os.path.exists(index_path): - webbrowser.open(f"file://{index_path}") - self.logger.info("Web interface opened in browser") - else: - self.logger.warning( - f"Web interface directory found but no index.html. " - f"Please check {web_dir} for setup instructions" - ) - else: - self.logger.info( - "Web interface not yet implemented. This feature will open a " - "web-based prompt management interface when the web_interface " - "directory is created." - ) - - except Exception as e: - self.logger.error(f"Error opening web interface: {e}") - def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]: """API method for JavaScript UI to search prompts.""" return self._search_prompts(search_text=search_text) diff --git a/prompt_manager_text.py b/prompt_manager_text.py index 37df856..7ba3123 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -53,7 +53,9 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): { "multiline": True, "dynamicPrompts": True, - "tooltip": "The text prompt to be processed and saved to database.", + "tooltip": ( + "The text prompt to be processed and saved to database." + ), }, ) }, @@ -62,35 +64,50 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): IO.STRING, { "default": "", - "tooltip": "Optional category for organizing prompts (e.g., 'landscapes', 'portraits')", + "tooltip": ( + "Optional category for organizing prompts " + "(e.g., 'landscapes', 'portraits')" + ), }, ), "tags": ( IO.STRING, { "default": "", - "tooltip": "Comma-separated tags for the prompt (e.g., 'anime, detailed, sunset')", + "tooltip": ( + "Comma-separated tags for the prompt " + "(e.g., 'anime, detailed, sunset')" + ), }, ), "search_text": ( IO.STRING, { "default": "", - "tooltip": 'Search past prompts by text; add tag:name (or tag:"two words") to filter by tag', + "tooltip": ( + "Search past prompts by text; add tag:name " + '(or tag:"two words") to filter by tag' + ), }, ), "prepend_text": ( IO.STRING, { "default": "", - "tooltip": "Text to prepend to the main prompt (connected STRING nodes will be added before the main text)", + "tooltip": ( + "Text to prepend to the main prompt (connected STRING " + "nodes will be added before the main text)" + ), }, ), "append_text": ( IO.STRING, { "default": "", - "tooltip": "Text to append to the main prompt (connected STRING nodes will be added after the main text)", + "tooltip": ( + "Text to append to the main prompt (connected STRING " + "nodes will be added after the main text)" + ), }, ), }, @@ -104,15 +121,17 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): RETURN_TYPES = (IO.STRING,) OUTPUT_TOOLTIPS = ( - "The final combined text string (with prepend/append applied) ready for use in other nodes.", + "The final combined text string (with prepend/append applied) ready for use " + "in other nodes.", ) FUNCTION = "process_text" OUTPUT_NODE = True CATEGORY = "🫢 ComfyAssets/🧠 Prompts" DESCRIPTION = ( - "Processes and manages text prompts with database storage and search capabilities. " - "Outputs a plain STRING that can be used with any node that accepts text input. " - "Includes all PromptManager features: categorization, tagging, search, and prepend/append functionality." + "Processes and manages text prompts with database storage and search " + "capabilities. Outputs a plain STRING that can be used with any node that " + "accepts text input. Includes all PromptManager features: categorization, " + "tagging, search, and prepend/append functionality." ) def process_text( @@ -153,7 +172,8 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): # Inject LoRA trigger words if integration is enabled final_text = self._inject_lora_trigger_words(final_text) - # For database storage, save the original main text with metadata about prepend/append + # For database storage, save the original main text with metadata + # about prepend/append storage_text = text # Save prompt to database and set execution context for gallery tracking @@ -223,10 +243,11 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): """ ComfyUI method to determine if node needs re-execution. - Returns a hash of input values that affect the text output. - This enables proper branch execution - only re-execute when inputs change. + Returns a hash of input values that affect the text output or the + saved prompt's metadata. This enables proper branch execution - only + re-execute when inputs change. """ import hashlib - combined = f"{text}|{prepend_text}|{append_text}" + combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}" return hashlib.sha256(combined.encode()).hexdigest() diff --git a/prompt_search_list.py b/prompt_search_list.py index 09a9319..eaab1fb 100644 --- a/prompt_search_list.py +++ b/prompt_search_list.py @@ -6,7 +6,6 @@ OUTPUT_IS_LIST=True, allowing direct connection to nodes that accept list inputs """ import re -from typing import Any, Dict, List, Tuple try: from .utils.logging_config import get_logger @@ -80,7 +79,9 @@ class PromptSearchList(ComfyNodeABC): IO.STRING, { "default": "", - "tooltip": "Comma-separated list of tags to filter by (partial match)", + "tooltip": ( + "Comma-separated list of tags to filter by (partial match)" + ), }, ), "min_rating": ( @@ -105,7 +106,10 @@ class PromptSearchList(ComfyNodeABC): "BOOLEAN", { "default": True, - "tooltip": "Skip prompts containing Clip_1/Clip_2/etc. multi-part markers", + "tooltip": ( + "Skip prompts containing Clip_1/Clip_2/etc. " + "multi-part markers" + ), }, ), }, diff --git a/py/api/__init__.py b/py/api/__init__.py index 90d4695..747fa30 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -13,7 +13,9 @@ import datetime import functools import gzip as gzip_module import json +import ntpath import os +import re from pathlib import Path from aiohttp import web @@ -29,6 +31,7 @@ from .lora_integration import LoraIntegrationMixin try: from ...database.operations import PromptDatabase from ...utils.logging_config import get_logger + from ...utils.video_metadata import is_video_path, read_video_metadata except ImportError: import sys @@ -37,6 +40,7 @@ except ImportError: ) from database.operations import PromptDatabase from utils.logging_config import get_logger + from utils.video_metadata import is_video_path, read_video_metadata def _get_project_root(): @@ -44,6 +48,147 @@ def _get_project_root(): return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +def _public_path(path) -> str: + """Render a server path for an API response without revealing the layout. + + Returns the path relative to the ComfyUI base directory (or to the parent + of an opted-in extra gallery root), as a POSIX-style string, or just the + basename when the path is not under any of those anchors. + """ + if not path: + return "" + from ..config import GalleryConfig + + try: + canonical = Path(os.path.normcase(os.path.realpath(str(path)))) + except (OSError, ValueError): + return os.path.basename(str(path)) + for anchor in GalleryConfig.path_anchors(): + if canonical.is_relative_to(Path(anchor)): + return canonical.relative_to(Path(anchor)).as_posix() or "." + return canonical.name + + +def _public_error(exc: BaseException) -> str: + """Error text safe for API responses: OSError paths reduced to basenames.""" + if isinstance(exc, OSError): + parts = [exc.strerror or type(exc).__name__] + if exc.filename: + parts.append(os.path.basename(str(exc.filename))) + return ": ".join(parts) + return str(exc) or type(exc).__name__ + + +# ── Static file serving ──────────────────────────────────────────────── +# Only these extensions are ever served from web/lib and web/js; the map +# doubles as the allow-list. +STATIC_MIME_TYPES = { + ".js": "application/javascript", + ".css": "text/css", + ".map": "application/json", + ".woff": "font/woff", + ".woff2": "font/woff2", + ".ttf": "font/ttf", + ".png": "image/png", + ".svg": "image/svg+xml", + ".ico": "image/x-icon", + ".html": "text/html", +} +_STATIC_SEPARATORS = re.compile(r"[\\/]+") + + +def _resolve_static_file(root, filepath): + """Real path of ``filepath`` under ``root``, or None when it must not be served. + + The request path is untrusted. It is rejected when it is empty, carries a + ``..`` segment, is absolute on any platform (a leading slash or backslash, + a drive letter such as ``C:`` or a UNC ``\\\\server\\share`` prefix all + re-anchor ``os.path.join`` on Windows), has an extension outside + STATIC_MIME_TYPES, or resolves (symlinks included) outside ``root``. + """ + if not isinstance(filepath, str) or not filepath: + return None + if ( + os.path.isabs(filepath) + or ntpath.isabs(filepath) + or ntpath.splitdrive(filepath)[0] + ): + return None + segments = _STATIC_SEPARATORS.split(filepath) + if any(segment in ("", ".", "..") for segment in segments): + return None + if os.path.splitext(filepath)[1].lower() not in STATIC_MIME_TYPES: + return None + try: + root_real = Path(os.path.normcase(os.path.realpath(root))) + candidate = Path( + os.path.normcase(os.path.realpath(os.path.join(root, *segments))) + ) + except (OSError, ValueError): + return None + if candidate == root_real or not candidate.is_relative_to(root_real): + return None + return str(candidate) + + +# ── JSON request bodies ──────────────────────────────────────────────── +JSON_BODY_MAX_BYTES = 1_000_000 +_JSON_READ_CHUNK = 64 * 1024 + + +def _json_error(message, status): + return web.json_response({"success": False, "error": message}, status=status) + + +async def _read_json_body(request, max_bytes=JSON_BODY_MAX_BYTES): + """Read and parse a JSON object body of at most ``max_bytes``. + + The body is read in chunks so an oversized or chunked upload is refused + as soon as it crosses the cap rather than buffered whole. + + Returns: + (data, None) on success, or (None, response) where ``response`` is a + 413 (body too large) or 400 (invalid JSON, or not a JSON object). + """ + too_large = _json_error(f"Request body too large (max {max_bytes} bytes)", 413) + declared = request.content_length + if declared is not None and declared > max_bytes: + return None, too_large + + chunks = [] + size = 0 + async for chunk in request.content.iter_chunked(_JSON_READ_CHUNK): + size += len(chunk) + if size > max_bytes: + return None, too_large + chunks.append(chunk) + + try: + data = json.loads(b"".join(chunks).decode("utf-8")) + except (UnicodeDecodeError, ValueError): + return None, _json_error("Request body must be valid JSON", 400) + if not isinstance(data, dict): + return None, _json_error("Request body must be a JSON object", 400) + return data, None + + +class _LongJobGuard: + """Async context manager handed out by ``_acquire_long_job``.""" + + def __init__(self, lock): + self._lock = lock + + async def __aenter__(self): + # The caller checked ``locked()`` on this same loop iteration, so this + # acquire never waits. + await self._lock.acquire() + return self + + async def __aexit__(self, exc_type, exc, tb): + self._lock.release() + return False + + # ── Gzip compression middleware ──────────────────────────────────────── _GZIP_MIN_SIZE = 1024 # Only compress bodies larger than 1 KB _GZIP_TYPES = frozenset( @@ -127,6 +272,14 @@ class PromptManagerAPI( self._html_cache = {} # Cached HTML file contents keyed by path self._gallery_cache = {} # dict: path_str -> (files, timestamp) self._gallery_cache_ttl = 30 # Cache TTL in seconds + # Single-flight locks for long jobs, keyed by job family; see + # _acquire_long_job(). Created lazily so no lock binds to an event + # loop before the server's loop exists. + self._long_jobs: dict = {} + # The current output scan (py/api/scan_job.ScanJob), if one was started. + self._scan_job = None + # The current duplicate scan (also a ScanJob), if one was started. + self._duplicate_job = None # Run cleanup on initialization to remove any existing duplicates try: @@ -140,6 +293,35 @@ class PromptManagerAPI( self.logger.info("PromptManager API initialization completed") + _public_path = staticmethod(_public_path) + _public_error = staticmethod(_public_error) + _read_json_body = staticmethod(_read_json_body) + + def _acquire_long_job(self, name): + """Claim the single-flight slot for a long job family, or None when busy. + + Usage in a handler:: + + job = self._acquire_long_job("duplicates") + if job is None: + return self._long_job_busy_response("Duplicate scan") + async with job: + ... + + Families are free-form strings ("duplicates", "thumbnails", + "autotag", ...); the lock is released when the ``async with`` block + exits, including on exceptions. Must be called from the event loop. + """ + lock = self._long_jobs.setdefault(name, asyncio.Lock()) + if lock.locked(): + return None + return _LongJobGuard(lock) + + @staticmethod + def _long_job_busy_response(label): + """409 envelope for a long job that is already running.""" + return _json_error(f"{label} already running", 409) + async def _run_in_executor(self, func, *args, **kwargs): """Run a blocking function in the default thread pool executor. @@ -181,191 +363,31 @@ class PromptManagerAPI( @routes.get("/prompt_manager/web") async def serve_web_ui(request): - try: - html_path = os.path.join(_get_project_root(), "web", "index.html") - - if os.path.exists(html_path): - with open(html_path, "r", encoding="utf-8") as f: - html_content = f.read() - - return web.Response( - text=html_content, content_type="text/html", charset="utf-8" - ) - else: - return web.Response( - text="

Web UI not found

HTML file not located at expected path.

", - content_type="text/html", - status=404, - ) - - except Exception as e: - self.logger.exception("Failed to load web UI") - return web.Response( - text="

Error

Failed to load web UI. Check server logs for details.

", - content_type="text/html", - status=500, - ) + return self._serve_html_page("index.html", "Web UI", cache=False) @routes.get("/prompt_manager/gallery.html") async def serve_gallery_ui(request): - try: - html_path = os.path.join( - _get_project_root(), - "web", - "metadata.html", - ) - - if html_path not in self._html_cache: - if os.path.exists(html_path): - with open(html_path, "r", encoding="utf-8") as f: - self._html_cache[html_path] = f.read() - else: - return web.Response( - text="

Gallery not found

gallery.html file not located at expected path.

", - content_type="text/html", - status=404, - ) - - return web.Response( - text=self._html_cache[html_path], - content_type="text/html", - charset="utf-8", - ) - - except Exception as e: - self.logger.exception("Failed to load gallery") - return web.Response( - text="

Error

Failed to load gallery. Check server logs for details.

", - content_type="text/html", - status=500, - ) + return self._serve_html_page("metadata.html", "Gallery") @routes.get("/prompt_manager/admin") async def serve_admin_ui(request): - try: - html_path = os.path.join( - _get_project_root(), - "web", - "admin.html", - ) - - if html_path not in self._html_cache: - if os.path.exists(html_path): - with open(html_path, "r", encoding="utf-8") as f: - self._html_cache[html_path] = f.read() - else: - return web.Response( - text="

Admin UI not found

", - content_type="text/html", - status=404, - ) - - return web.Response( - text=self._html_cache[html_path], - content_type="text/html", - charset="utf-8", - ) - - except Exception as e: - self.logger.exception("Failed to load admin UI") - return web.Response( - text="

Error

Failed to load admin UI. Check server logs for details.

", - content_type="text/html", - status=500, - ) + return self._serve_html_page("admin.html", "Admin UI") @routes.get("/prompt_manager/gallery") async def serve_gallery_admin_ui(request): - try: - html_path = os.path.join( - _get_project_root(), - "web", - "gallery.html", - ) - - if html_path not in self._html_cache: - if os.path.exists(html_path): - with open(html_path, "r", encoding="utf-8") as f: - self._html_cache[html_path] = f.read() - else: - return web.Response( - text="

Gallery not found

gallery.html file not located at expected path.

", - content_type="text/html", - status=404, - ) - - return web.Response( - text=self._html_cache[html_path], - content_type="text/html", - charset="utf-8", - ) - - except Exception as e: - self.logger.exception("Failed to load gallery") - return web.Response( - text="

Error

Failed to load gallery. Check server logs for details.

", - content_type="text/html", - status=500, - ) + return self._serve_html_page("gallery.html", "Gallery") # ── Static file serving ─────────────────────────────────────── @routes.get("/prompt_manager/lib/{filepath:.*}") async def serve_lib_static(request): """Serve static library files (JS, CSS) from web/lib directory.""" - MIME_TYPES = { - ".js": "application/javascript", - ".css": "text/css", - ".json": "application/json", - ".map": "application/json", - } - - filepath = request.match_info.get("filepath", "") - - # Security: prevent directory traversal - if ".." in filepath or filepath.startswith("/"): - return web.Response(text="Forbidden", status=403) - - file_path = os.path.join(_get_project_root(), "web", "lib", filepath) - - if not os.path.exists(file_path) or not os.path.isfile(file_path): - return web.Response(text=f"Not Found: {filepath}", status=404) - - ext = os.path.splitext(file_path)[1].lower() - content_type = MIME_TYPES.get(ext, "application/octet-stream") - - with open(file_path, "rb") as f: - content = f.read() - - return web.Response(body=content, content_type=content_type) + return self._serve_static_file("lib", request.match_info.get("filepath")) @routes.get("/prompt_manager/js/{filepath:.*}") async def serve_js_static(request): """Serve static JavaScript files from web/js directory.""" - MIME_TYPES = { - ".js": "application/javascript", - ".css": "text/css", - ".json": "application/json", - ".map": "application/json", - } - - filepath = request.match_info.get("filepath", "") - - if ".." in filepath or filepath.startswith("/"): - return web.Response(text="Forbidden", status=403) - - file_path = os.path.join(_get_project_root(), "web", "js", filepath) - - if not os.path.exists(file_path) or not os.path.isfile(file_path): - return web.Response(text=f"Not Found: {filepath}", status=404) - - ext = os.path.splitext(file_path)[1].lower() - content_type = MIME_TYPES.get(ext, "application/octet-stream") - - with open(file_path, "rb") as f: - content = f.read() - - return web.Response(body=content, content_type=content_type) + return self._serve_static_file("js", request.match_info.get("filepath")) # ── Register domain-specific routes from mixins ─────────────── @@ -390,6 +412,47 @@ class PromptManagerAPI( self.logger.info("All routes registered with decorator pattern") + def _serve_static_file(self, subdir, filepath): + """Serve ``web//``; 403 outside it, 404 when missing.""" + root = os.path.join(_get_project_root(), "web", subdir) + file_path = _resolve_static_file(root, filepath) + if file_path is None: + return web.Response(text="Forbidden", status=403) + if not os.path.isfile(file_path): + return web.Response(text="Not Found", status=404) + ext = os.path.splitext(file_path)[1].lower() + with open(file_path, "rb") as f: + content = f.read() + return web.Response(body=content, content_type=STATIC_MIME_TYPES[ext]) + + def _serve_html_page(self, filename, title, cache=True): + """Serve ``web/``; 404 when missing, 500 when unreadable.""" + html_path = os.path.join(_get_project_root(), "web", filename) + try: + if cache and html_path in self._html_cache: + content = self._html_cache[html_path] + elif os.path.exists(html_path): + with open(html_path, "r", encoding="utf-8") as f: + content = f.read() + if cache: + self._html_cache[html_path] = content + else: + return web.Response( + text=f"

{title} not found

" + "

HTML file not located at expected path.

", + content_type="text/html", + status=404, + ) + return web.Response(text=content, content_type="text/html", charset="utf-8") + except Exception: + self.logger.exception(f"Failed to load {title}") + return web.Response( + text=f"

Error

Failed to load {title}. " + "Check server logs for details.

", + content_type="text/html", + status=500, + ) + # ── Shared utilities used by multiple mixins ────────────────────── def _enrich_images(self, images): @@ -415,7 +478,8 @@ class PromptManagerAPI( rel_path = img_path.resolve().relative_to(output_path.resolve()) image["relative_path"] = str(rel_path) image["url"] = ( - f"/prompt_manager/images/serve/{url_quote(rel_path.as_posix(), safe='/')}" + "/prompt_manager/images/serve/" + f"{url_quote(rel_path.as_posix(), safe='/')}" ) # Check for thumbnail @@ -423,11 +487,13 @@ class PromptManagerAPI( thumb_rel = ( f"thumbnails/{rel_no_ext.as_posix()}_thumb{rel_path.suffix}" ) - thumb_abs = output_path / thumb_rel - if thumb_abs.exists(): - image["thumbnail_url"] = ( - f"/prompt_manager/images/serve/{url_quote(thumb_rel, safe='/')}" - ) + from .prompts import thumbnail_url_for + + thumbnail_url = thumbnail_url_for( + thumb_rel, output_path / thumb_rel + ) + if thumbnail_url: + image["thumbnail_url"] = thumbnail_url break # Found matching root, stop searching except (ValueError, RuntimeError): continue # Try next root @@ -476,17 +542,18 @@ class PromptManagerAPI( from ..config import GalleryConfig if GalleryConfig.MONITORING_DIRECTORIES: - configured_dir = Path(GalleryConfig.MONITORING_DIRECTORIES[0]).resolve() - if configured_dir.exists() and configured_dir.is_dir(): + configured = GalleryConfig.MONITORING_DIRECTORIES[0] + # Same check as the settings endpoint: a root hand-edited + # into config.json must not open a directory outside ComfyUI + ok, reason = GalleryConfig.validate_gallery_root(configured) + if ok: + configured_dir = Path(configured).resolve() self.logger.info( f"Using configured monitoring directory: {configured_dir}" ) self._cached_output_dir = str(configured_dir) return self._cached_output_dir - else: - self.logger.warning( - f"Configured directory does not exist: {GalleryConfig.MONITORING_DIRECTORIES[0]}" - ) + self.logger.warning(f"Ignoring configured directory: {reason}") except ImportError: self.logger.debug( "GalleryConfig not available, skipping configured directory check" @@ -506,7 +573,8 @@ class PromptManagerAPI( output_dir = current_dir / "output" if output_dir.exists() and output_dir.is_dir(): self.logger.debug( - f"Found ComfyUI output directory via upward search: {output_dir}" + "Found ComfyUI output directory via upward search: " + f"{output_dir}" ) self._cached_output_dir = str(output_dir) return self._cached_output_dir @@ -557,39 +625,72 @@ class PromptManagerAPI( return None def _get_all_output_dirs(self): - """Get all configured output directories, falling back to auto-detect. + """Get the directories images may be served from. + + Configured gallery roots are filtered through + ``GalleryConfig.validate_gallery_root`` so a root that was hand-edited + into config.json (or persisted before validation existed) can never + expose files outside ComfyUI's directories. Returns: - List[Path]: Valid output directory paths, possibly empty. + List[Path]: Validated directory paths, possibly empty. """ from ..config import GalleryConfig + configured = list(GalleryConfig.MONITORING_DIRECTORIES) output_dirs = [] - if GalleryConfig.MONITORING_DIRECTORIES: - for d in GalleryConfig.MONITORING_DIRECTORIES: - p = Path(d).resolve() - if p.is_dir(): - output_dirs.append(p) + seen = set() + for d in configured: + ok, reason = GalleryConfig.validate_gallery_root(d) + if not ok: + self.logger.warning(f"Ignoring configured gallery root: {reason}") + continue + real = os.path.realpath(d) + key = os.path.normcase(real) # one entry per directory, however spelled + if key in seen: + continue + seen.add(key) + output_dirs.append(Path(real)) - # Fallback to auto-detect if no configured dirs are valid - if not output_dirs: - fallback = self._find_comfyui_output_dir() - if fallback: - output_dirs.append(Path(fallback)) + if output_dirs: + return output_dirs - return output_dirs + if configured: + # Every configured root was rejected: never let auto-detection + # resurrect it, serve ComfyUI's own output directory instead. + return self._comfyui_output_dir_list() + + fallback = self._find_comfyui_output_dir() + return [Path(fallback)] if fallback else [] + + def _comfyui_output_dir_list(self): + """ComfyUI's output directory as a one-element list, or [] if unknown.""" + try: + import folder_paths + + output_dir = folder_paths.get_output_directory() + except (ImportError, AttributeError, OSError): + return [] + if isinstance(output_dir, str) and os.path.isdir(output_dir): + return [Path(os.path.realpath(output_dir))] + return [] def _extract_comfyui_metadata(self, image_path): - """Extract ComfyUI workflow metadata from PNG image files.""" + """Embedded ComfyUI metadata of an image (PNG text chunks) or a video. + + Videos go through ffprobe (see utils/video_metadata). Unreadable files + are a warning, not an error: a scan over a large output tree meets + truncated and foreign files, and it carries on past them. + """ + if is_video_path(image_path): + return read_video_metadata(image_path) try: with Image.open(image_path) as img: - metadata = {} - if hasattr(img, "text"): - for key, value in img.text.items(): - metadata[key] = value - return metadata + return dict(getattr(img, "text", None) or {}) except Exception as e: - self.logger.error(f"Error reading {image_path}: {e}") + self.logger.warning( + f"Unreadable image {os.path.basename(str(image_path))}: {e}" + ) return {} def _parse_comfyui_prompt(self, metadata): @@ -710,7 +811,8 @@ class PromptManagerAPI( New format: inputs is a list of connection objects inputs = [ - {"name": "text", "type": "STRING", "link": null, "widget": {"name": "text"}}, + {"name": "text", "type": "STRING", "link": null, + "widget": {"name": "text"}}, {"name": "clip", "type": "CLIP", "link": 11} ] """ @@ -773,7 +875,10 @@ class PromptManagerAPI( return None def _extract_positive_prompt_from_comfyui_data(self, data): - """Extract positive prompt from ComfyUI data, handling both old and new formats.""" + """Extract positive prompt from ComfyUI data. + + Handles both old and new formats. + """ if not isinstance(data, dict): return None diff --git a/py/api/admin.py b/py/api/admin.py index ae9b6a7..d13c034 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -1,22 +1,302 @@ """Admin and maintenance API routes for PromptManager.""" import asyncio +import contextlib +from contextlib import closing import datetime import hashlib import json import os -import shutil import sqlite3 import tempfile -import traceback +import threading from pathlib import Path from aiohttp import web +from .prompts import thumbnail_url_for +from .scan_job import ScanJob + try: - from ...utils.validators import validate_result_timeout + from ...utils.parallel import map_parallel + from ...utils.validators import validate_result_timeout, validate_worker_threads except ImportError: - from utils.validators import validate_result_timeout + from utils.parallel import map_parallel + from utils.validators import validate_result_timeout, validate_worker_threads + +try: + from ...database.operations import PromptDatabase +except ImportError: + from database.operations import PromptDatabase + + +try: + from ...utils.hashing import generate_prompt_hash +except ImportError: + from utils.hashing import generate_prompt_hash + +IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff") +VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv") +SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan +MAX_SAFETY_BACKUPS = 10 # '.backup_*' files kept next to the live database +MAX_SCAN_FILES = 50_000 # Upper bound on media files one scan will look at +MAX_SCAN_DEPTH = 12 # Directory levels below an output root that are scanned +# Files the duplicate cleaner may delete: everything the scan reports, plus +# the common single-f TIFF spelling. Anything else inside the output +# directory (databases, scripts, notes) is never a "duplicate image". +DELETABLE_MEDIA_EXTENSIONS = frozenset(IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + (".tif",)) + + +def _collect_media_files(output_dirs, extensions): + """Media files under ``output_dirs`` (blocking), bounded and symlink-free. + + Walks each root without following symlinks, skips ``thumbnails`` + directories and symlinked files, never descends more than MAX_SCAN_DEPTH + levels and stops after MAX_SCAN_FILES files in total. Extensions match + case-insensitively and each path is listed once. + """ + wanted = {ext.lower() for ext in extensions} + found = [] + seen = set() + for output_dir in output_dirs: + root = os.path.abspath(str(output_dir)) + if not os.path.isdir(root): + continue + for current, dirnames, filenames in os.walk(root, followlinks=False): + depth = 0 if current == root else len(Path(current).relative_to(root).parts) + dirnames[:] = sorted( + d + for d in dirnames + if d != "thumbnails" + and depth < MAX_SCAN_DEPTH + and not os.path.islink(os.path.join(current, d)) + ) + for name in sorted(filenames): + if os.path.splitext(name)[1].lower() not in wanted: + continue + full = os.path.join(current, name) + if os.path.islink(full) or not os.path.isfile(full): + continue + key = os.path.normcase(full) + if key in seen: + continue + seen.add(key) + found.append(Path(full)) + if len(found) >= MAX_SCAN_FILES: + return found + return found + + +def _thumbnail_rel_path(rel_path, thumbnail_ext): + """Relative path of the thumbnail the gallery would generate for ``rel_path``.""" + return f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb{thumbnail_ext}" + + +async def _read_json(request): + """Parse a JSON object body through the package's capped reader. + + Returns: + (data, None) on success, or (None, response) carrying a 400/413 reply. + """ + from . import _read_json_body + + return await _read_json_body(request) + + +def _sse(payload): + """Encode one server-sent event.""" + return f"data: {json.dumps(payload)}\n\n" + + +def _inspect_restore_upload(path): + """Checks on a restore upload beyond PromptModel.verify_database_file. + + Opens the file read-only with ``PRAGMA trusted_schema=OFF`` (so nothing + in the schema runs while it is inspected) and refuses databases that + carry triggers or views: the prompts database never needs either, and a + crafted upload could use them to run SQL against the live data later. + + Returns: + None when acceptable, otherwise a human-readable reason (blocking). + """ + uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro" + try: + conn = sqlite3.connect(uri, uri=True) + except sqlite3.Error as e: + return f"Invalid SQLite database: {e}" + try: + conn.execute("PRAGMA trusted_schema=OFF") + rows = conn.execute( + "SELECT type, name FROM sqlite_master " + "WHERE type IN ('trigger', 'view') ORDER BY type, name" + ).fetchall() + except sqlite3.Error as e: + return f"Invalid SQLite database: {e}" + finally: + conn.close() + if rows: + listed = ", ".join(f"{kind} {name}" for kind, name in rows[:5]) + return ( + "Database contains triggers or views, which a prompts database " + f"never has: {listed}" + ) + return None + + +def _prune_safety_backups(db_path, keep=None): + """Delete the oldest ``.backup_*`` files beyond ``keep`` (blocking). + + Returns the paths removed. Errors deleting a file are ignored so a + stubborn old backup never fails the restore that just succeeded. + """ + if keep is None: + keep = MAX_SAFETY_BACKUPS + db_path = os.path.abspath(db_path) + directory = os.path.dirname(db_path) + prefix = os.path.basename(db_path) + ".backup_" + try: + names = os.listdir(directory) + except OSError: + return [] + candidates = [ + os.path.join(directory, name) + for name in names + if name.startswith(prefix) and os.path.isfile(os.path.join(directory, name)) + ] + candidates.sort(key=lambda p: (os.path.getmtime(p), p)) + excess = candidates[: max(0, len(candidates) - keep)] + removed = [] + for path in excess: + try: + os.remove(path) + removed.append(path) + except OSError: + continue + return removed + + +def _public_helpers(): + """The package's public-path helpers (imported lazily to avoid a cycle).""" + from . import _public_error, _public_path + + return _public_path, _public_error + + +def _diagnose_database(db_path): + """Row counts straight from the database file (small blocking query).""" + _public_path, _public_error = _public_helpers() + if not os.path.exists(db_path): + return { + "status": "error", + "message": f"Database file not found: {os.path.basename(db_path)}", + } + try: + with closing(sqlite3.connect(db_path)) as conn: + prompt_count = conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0] + has_images_table = ( + conn.execute( + "SELECT name FROM sqlite_master " + "WHERE type='table' AND name='generated_images'" + ).fetchone() + is not None + ) + image_count = 0 + if has_images_table: + image_count = conn.execute( + "SELECT COUNT(*) FROM generated_images" + ).fetchone()[0] + return { + "status": "ok", + "prompt_count": prompt_count, + "has_images_table": has_images_table, + "image_count": image_count, + } + except Exception as e: + return {"status": "error", "message": f"Database error: {_public_error(e)}"} + + +def _diagnose_dependencies(): + """Presence of the optional runtime dependencies.""" + dependencies = {"sqlite3": True} + for name in ("watchdog", "PIL"): + try: + __import__(name) + dependencies[name] = True + except ImportError: + dependencies[name] = False + return { + "status": "ok" if all(dependencies.values()) else "error", + "dependencies": dependencies, + } + + +def _diagnose_output_dirs(configured_dirs): + """Configured gallery roots, ComfyUI's output dir, or relative fallbacks.""" + _public_path, _ = _public_helpers() + candidates = list(configured_dirs) + try: + import folder_paths + + candidates.append(folder_paths.get_output_directory()) + except ImportError: + pass + + output_dirs = [] + seen = set() + + def add(path): + # One entry per directory: resolve links and fold case (Windows) + real = os.path.realpath(path) + key = os.path.normcase(real) + if os.path.exists(real) and key not in seen: + seen.add(key) + output_dirs.append(real) + + for candidate in candidates: + if candidate: + add(candidate) + + if not output_dirs: + for rel in ("output", "../output", "../../output"): + add(rel) + + return { + "status": "ok" if output_dirs else "warning", + "output_dirs": [_public_path(d) for d in output_dirs], + } + + +def _diagnose_image_monitor(): + """Status of the running image monitor, with public directory names.""" + _public_path, _public_error = _public_helpers() + try: + monitor = _current_image_monitor() + if monitor is None: + return {"status": "error", "message": "Image monitor not initialized"} + status = dict(monitor.get_status()) + status["monitored_directories"] = [ + _public_path(d) for d in status.get("monitored_directories", []) + ] + return {"status": "ok" if status.get("observer_alive") else "error", **status} + except Exception as e: + return { + "status": "error", + "message": f"Failed to get monitor status: {_public_error(e)}", + } + + +def _image_monitor_module(): + """The utils.image_monitor module under either package identity.""" + try: + from ...utils import image_monitor + except ImportError: + from utils import image_monitor + return image_monitor + + +def _current_image_monitor(): + """The running ImageMonitor singleton, or None when not started.""" + return getattr(_image_monitor_module(), "_monitor_instance", None) class AdminRoutesMixin: @@ -27,6 +307,14 @@ class AdminRoutesMixin: async def scan_duplicates_route(request): return await self.scan_duplicates_endpoint(request) + @routes.post("/prompt_manager/scan_duplicates/stream") + async def scan_duplicates_stream_route(request): + return await self.scan_duplicates_stream(request) + + @routes.get("/prompt_manager/scan_duplicates/status") + async def scan_duplicates_status_route(request): + return await self.scan_duplicates_status(request) + @routes.post("/prompt_manager/delete_duplicate_images") async def delete_duplicate_images_route(request): return await self.delete_duplicate_images_endpoint(request) @@ -71,8 +359,21 @@ class AdminRoutesMixin: async def scan_images_route(request): return await self.scan_images(request) + @routes.get("/prompt_manager/scan/status") + async def scan_status_route(request): + return await self.scan_status(request) + + LONG_JOB_DUPLICATES = "duplicates" + async def scan_duplicates_endpoint(self, request): - """Scan for duplicate images without removing them.""" + """Scan for duplicate images without removing them (single flight).""" + job = self._acquire_long_job(self.LONG_JOB_DUPLICATES) + if job is None: + return self._long_job_busy_response("Duplicate scan") + async with job: + return await self._scan_duplicates(request) + + async def _scan_duplicates(self, request): try: duplicates = await self.find_duplicate_images() @@ -85,15 +386,95 @@ class AdminRoutesMixin: ) except Exception as e: - self.logger.error(f"Scan duplicates error: {e}") + self.logger.error(f"Scan duplicates error: {e}", exc_info=True) return web.json_response( - { - "success": False, - "error": f"Failed to scan duplicate images: {str(e)}", - }, + {"success": False, "error": "Failed to scan duplicate images"}, status=500, ) + async def scan_duplicates_stream(self, request): + """Run the duplicate scan as a background job and stream its progress. + + A request while a scan is running attaches to it. The job holds the + duplicates lock, so a delete is refused (409) until it finishes. + """ + job = self._duplicate_job + if job is None or job.finished: + guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES) + if guard is None: + return self._long_job_busy_response("Duplicate scan or delete") + job = ScanJob() + self._duplicate_job = job + job.task = asyncio.ensure_future(self._run_duplicate_job(job, guard)) + return await self._stream_scan_job(request, job) + + async def scan_duplicates_status(self, request): + """Whether a duplicate scan is running, and its latest event.""" + job = self._duplicate_job + return web.json_response( + { + "success": True, + "running": bool(job and job.running), + "last_event": job.last_event if job else None, + } + ) + + async def _run_duplicate_job(self, job, guard): + loop = asyncio.get_running_loop() + counts = {"processed": 0} + + def report(done, total): + # Called from worker threads; hand the event to the loop. + counts["processed"] = done + pct = int(done / total * 100) if total else 100 + loop.call_soon_threadsafe( + job.publish, + { + "type": "progress", + "progress": pct, + "status": f"Hashing file {done}/{total}...", + "processed": done, + "found": 0, + }, + ) + + try: + async with guard: + job.publish( + { + "type": "progress", + "progress": 0, + "status": "Collecting media files...", + "processed": 0, + "found": 0, + } + ) + output_dir = await self._run_in_executor(self._find_comfyui_output_dir) + groups = await self._run_in_executor( + self._find_duplicate_images_sync, output_dir, report + ) + job.publish( + { + "type": "complete", + "processed": counts["processed"], + "found": len(groups), + "duplicates": groups, + "message": f"Found {len(groups)} groups of duplicate images", + } + ) + except Exception: + self.logger.exception("Duplicate scan error") + job.publish( + { + "type": "error", + "message": ( + "An internal error occurred. Check server logs for details." + ), + } + ) + finally: + job.finish() + async def cleanup_duplicates_endpoint(self, request): """Cleanup duplicate prompts endpoint.""" try: @@ -108,122 +489,99 @@ class AdminRoutesMixin: ) except Exception as e: - self.logger.error(f"Cleanup error: {e}") + self.logger.error(f"Cleanup error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to cleanup duplicates: {str(e)}"}, + {"success": False, "error": "Failed to cleanup duplicates"}, status=500, ) async def find_duplicate_images(self): - """Find duplicate images in ComfyUI output directory using content hashing.""" - self.logger.info("Scanning for duplicate images") + """Find duplicate media in the output directory, off the event loop.""" + output_dir = await self._run_in_executor(self._find_comfyui_output_dir) + return await self._run_in_executor(self._find_duplicate_images_sync, output_dir) - try: - output_dir = self._find_comfyui_output_dir() - if not output_dir: - self.logger.warning("ComfyUI output directory not found") - return [] + def _find_duplicate_images_sync(self, output_dir, progress=None): + """Group media files under ``output_dir`` by content hash (blocking). - output_path = Path(output_dir) + Files are hashed on ``PromptManagerConfig.WORKER_THREADS`` threads. + ``progress(done, total)`` is called from those threads about every + percent and once at the end. + """ + from ..config import PromptManagerConfig - image_extensions = [ - ".png", - ".jpg", - ".jpeg", - ".gif", - ".webp", - ".bmp", - ".tiff", - ] - video_extensions = [".mp4", ".avi", ".mov", ".mkv", ".webm", ".gif"] - all_extensions = image_extensions + video_extensions - - media_files = [] - seen_paths = set() - for ext in all_extensions: - for pattern in [f"*{ext.lower()}", f"*{ext.upper()}"]: - for media_path in output_path.rglob(pattern): - if "thumbnails" not in media_path.parts: - normalized_path = str(media_path).lower() - if normalized_path not in seen_paths: - seen_paths.add(normalized_path) - media_files.append(media_path) - - self.logger.info(f"Found {len(media_files)} media files to analyze") - - file_hashes = {} - processed = 0 - - for media_path in media_files: - try: - file_hash = self._calculate_file_hash(media_path) - - if file_hash not in file_hashes: - file_hashes[file_hash] = [] - - stat = media_path.stat() - rel_path = media_path.relative_to(output_path) - extension = media_path.suffix.lower() - is_video = extension in [ext.lower() for ext in video_extensions] - media_type = "video" if is_video else "image" - - # Check if thumbnail exists - thumbnail_url = None - thumbnails_dir = output_path / "thumbnails" - if thumbnails_dir.exists(): - thumbnail_ext = ".jpg" if is_video else extension - rel_path_no_ext = rel_path.with_suffix("") - thumbnail_rel_path = f"thumbnails/{rel_path_no_ext.as_posix()}_thumb{thumbnail_ext}" - thumbnail_abs_path = output_path / thumbnail_rel_path - - if thumbnail_abs_path.exists(): - from urllib.parse import quote - - thumbnail_url = f'/prompt_manager/images/serve/{quote(thumbnail_rel_path, safe="/")}' - - image_info = { - "id": str(hash(str(media_path))), - "filename": media_path.name, - "path": str(media_path), - "relative_path": str(rel_path), - "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", - "thumbnail_url": thumbnail_url, - "size": stat.st_size, - "modified_time": stat.st_mtime, - "extension": extension, - "media_type": media_type, - "is_video": is_video, - "hash": file_hash, - } - - file_hashes[file_hash].append(image_info) - processed += 1 - - if processed % 100 == 0: - self.logger.info( - f"Processed {processed}/{len(media_files)} files for duplicate detection" - ) - - except Exception as e: - self.logger.error(f"Error processing file {media_path}: {e}") - continue - - # Find duplicates (groups with more than one file) - duplicates = [] - for file_hash, images in file_hashes.items(): - if len(images) > 1: - images.sort(key=lambda x: x["modified_time"]) - duplicates.append( - {"hash": file_hash, "images": images, "count": len(images)} - ) - - self.logger.info(f"Found {len(duplicates)} groups of duplicate images") - return duplicates - - except Exception as e: - self.logger.error(f"Error finding duplicate images: {e}") + if not output_dir or not os.path.isdir(output_dir): + self.logger.warning("ComfyUI output directory not found") return [] + output_path = Path(output_dir) + extensions = IMAGE_EXTENSIONS + VIDEO_EXTENSIONS + media_files = _collect_media_files([output_path], extensions) + total = len(media_files) + self.logger.info(f"Found {total} media files to analyze") + step = max(1, total // 100) + counter = {"done": 0} + lock = threading.Lock() + + def describe(media_path): + try: + info = self._describe_media_file(media_path, output_path) + except Exception as e: + self.logger.error(f"Error processing file {media_path.name}: {e}") + info = None + with lock: + counter["done"] += 1 + done = counter["done"] + if done % 100 == 0: + self.logger.info( + f"Processed {done}/{total} files " "for duplicate detection" + ) + if progress and (done % step == 0 or done == total): + progress(done, total) + return info + + infos = map_parallel(describe, media_files, PromptManagerConfig.WORKER_THREADS) + + file_hashes = {} + for info in infos: + if info is not None: + file_hashes.setdefault(info["hash"], []).append(info) + + duplicates = [] + for file_hash, images in file_hashes.items(): + if len(images) > 1: + images.sort(key=lambda x: x["modified_time"]) + duplicates.append( + {"hash": file_hash, "images": images, "count": len(images)} + ) + + self.logger.info(f"Found {len(duplicates)} groups of duplicate images") + return duplicates + + def _describe_media_file(self, media_path, output_path): + """Hash one media file and describe it for the duplicates response.""" + stat = media_path.stat() + rel_path = media_path.relative_to(output_path) + extension = media_path.suffix.lower() + is_video = extension in VIDEO_EXTENSIONS + + thumb_rel = _thumbnail_rel_path(rel_path, ".jpg" if is_video else extension) + thumbnail_url = thumbnail_url_for(thumb_rel, output_path / thumb_rel) + + return { + "id": str(hash(str(media_path))), + "filename": media_path.name, + "path": str(rel_path), + "relative_path": str(rel_path), + "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", + "thumbnail_url": thumbnail_url, + "size": stat.st_size, + "modified_time": stat.st_mtime, + "extension": extension, + "media_type": "video" if is_video else "image", + "is_video": is_video, + "hash": self._calculate_file_hash(media_path), + } + def _calculate_file_hash(self, file_path): """Calculate SHA-256 hash of a file's content.""" hash_sha256 = hashlib.sha256() @@ -233,101 +591,108 @@ class AdminRoutesMixin: return hash_sha256.hexdigest() async def delete_duplicate_images_endpoint(self, request): - """Delete duplicate image files from disk.""" + """Delete duplicate image files from disk (single flight with the scan).""" + job = self._acquire_long_job(self.LONG_JOB_DUPLICATES) + if job is None: + return self._long_job_busy_response("Duplicate scan or delete") + async with job: + return await self._delete_duplicate_images(request) + + async def _delete_duplicate_images(self, request): try: - data = await request.json() + data, error_response = await _read_json(request) + if error_response is not None: + return error_response image_paths = data.get("image_paths", []) - if not image_paths: + if not image_paths or not isinstance(image_paths, list): return web.json_response( {"success": False, "error": "No image paths provided"}, status=400, ) - deleted_count = 0 - failed_count = 0 - failed_files = [] - - for image_path in image_paths: - try: - # Ensure the path is within the output directory for security - output_dir = self._find_comfyui_output_dir() - if not output_dir: - failed_files.append( - f"{image_path} (output directory not found)" - ) - failed_count += 1 - continue - - output_path = Path(output_dir) - file_path = Path(image_path) - - # Security check - ensure file is within output directory - try: - file_path.resolve().relative_to(output_path.resolve()) - except ValueError: - self.logger.warning( - f"Attempted to delete file outside output directory: {image_path}" - ) - failed_files.append(f"{image_path} (outside output directory)") - failed_count += 1 - continue - - if file_path.exists() and file_path.is_file(): - os.remove(file_path) - deleted_count += 1 - self.logger.info(f"Deleted duplicate image: {image_path}") - - # Also try to delete associated thumbnail if it exists - try: - rel_path = file_path.relative_to(output_path) - rel_path_no_ext = rel_path.with_suffix("") - thumbnail_path = ( - output_path - / "thumbnails" - / f"{rel_path_no_ext.as_posix()}_thumb{file_path.suffix}" - ) - if thumbnail_path.exists(): - os.remove(thumbnail_path) - self.logger.debug( - f"Deleted associated thumbnail: {thumbnail_path}" - ) - except Exception as e: - self.logger.warning( - f"Could not delete thumbnail for {image_path}: {e}" - ) - else: - failed_files.append(f"{image_path} (file not found)") - failed_count += 1 - - except Exception as e: - self.logger.error(f"Error deleting file {image_path}: {e}") - failed_files.append(f"{image_path} ({str(e)})") - failed_count += 1 + output_dir = await self._run_in_executor(self._find_comfyui_output_dir) + result = await self._run_in_executor( + self._delete_duplicate_images_sync, image_paths, output_dir + ) response_data = { "success": True, - "deleted_count": deleted_count, - "failed_count": failed_count, - "message": f"Deleted {deleted_count} files successfully", + "deleted_count": result["deleted_count"], + "failed_count": result["failed_count"], + "message": f"Deleted {result['deleted_count']} files successfully", } - - if failed_count > 0: - response_data["failed_files"] = failed_files - response_data["message"] += f", {failed_count} failed" + if result["failed_count"] > 0: + response_data["failed_files"] = result["failed_files"] + response_data["message"] += f", {result['failed_count']} failed" return web.json_response(response_data) except Exception as e: - self.logger.error(f"Delete duplicate images error: {e}") + self.logger.error(f"Delete duplicate images error: {e}", exc_info=True) return web.json_response( - { - "success": False, - "error": f"Failed to delete duplicate images: {str(e)}", - }, + {"success": False, "error": "Failed to delete duplicate images"}, status=500, ) + def _delete_duplicate_images_sync(self, image_paths, output_dir): + """Delete the given files (absolute or output-relative) within + ``output_dir``.""" + result = {"deleted_count": 0, "failed_count": 0, "failed_files": []} + output_path = Path(output_dir) if output_dir else None + + for image_path in image_paths: + if output_path is None: + failure = "output directory not found" + else: + try: + failure = self._delete_one_duplicate(str(image_path), output_path) + except Exception as e: + self.logger.error(f"Error deleting file {image_path}: {e}") + failure = self._public_error(e) + + if failure is None: + result["deleted_count"] += 1 + else: + result["failed_count"] += 1 + result["failed_files"].append(f"{image_path} ({failure})") + + return result + + def _delete_one_duplicate(self, image_path, output_path): + """Delete one file and its thumbnail; return a failure reason or None.""" + file_path = Path(image_path) + if not file_path.is_absolute(): + file_path = output_path / file_path + + try: + rel_path = file_path.resolve().relative_to(output_path.resolve()) + except ValueError: + self.logger.warning( + f"Attempted to delete file outside output directory: {image_path}" + ) + return "outside output directory" + + if file_path.suffix.lower() not in DELETABLE_MEDIA_EXTENSIONS: + return "not a media file" + if "thumbnails" in rel_path.parts: + return "thumbnails are managed by the gallery" + + if not file_path.is_file(): + return "file not found" + + os.remove(file_path) + self.logger.info(f"Deleted duplicate image: {image_path}") + + thumbnail_path = output_path / _thumbnail_rel_path(rel_path, file_path.suffix) + try: + if thumbnail_path.exists(): + os.remove(thumbnail_path) + self.logger.debug(f"Deleted associated thumbnail: {thumbnail_path}") + except OSError as e: + self.logger.warning(f"Could not delete thumbnail for {image_path}: {e}") + return None + async def get_statistics(self, request): """Get database statistics.""" try: @@ -336,9 +701,9 @@ class AdminRoutesMixin: return web.json_response({"success": True, "stats": stats}) except Exception as e: - self.logger.error(f"Stats error: {e}") + self.logger.error(f"Stats error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to get statistics: {str(e)}"}, + {"success": False, "error": "Failed to get statistics"}, status=500, ) @@ -347,192 +712,69 @@ class AdminRoutesMixin: try: from ..config import PromptManagerConfig, GalleryConfig - # Get monitored directories from image monitor singleton if available - monitored_dirs = [] - try: - import sys - - monitor_module = None - for mod_name in list(sys.modules.keys()): - if "image_monitor" in mod_name and hasattr( - sys.modules[mod_name], "_monitor_instance" - ): - monitor_module = sys.modules[mod_name] - break - - if monitor_module and monitor_module._monitor_instance is not None: - monitored_dirs = getattr( - monitor_module._monitor_instance, "monitored_directories", [] - ) - elif GalleryConfig.MONITORING_DIRECTORIES: - monitored_dirs = GalleryConfig.MONITORING_DIRECTORIES - except Exception: - if GalleryConfig.MONITORING_DIRECTORIES: - monitored_dirs = GalleryConfig.MONITORING_DIRECTORIES - + root_paths = [ + self._public_path(d) for d in GalleryConfig.MONITORING_DIRECTORIES + ] return web.json_response( { "success": True, "settings": { "result_timeout": PromptManagerConfig.RESULT_TIMEOUT, "webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE, - "gallery_root_paths": list( - GalleryConfig.MONITORING_DIRECTORIES - ), - "gallery_root_path": ( - GalleryConfig.MONITORING_DIRECTORIES[0] - if GalleryConfig.MONITORING_DIRECTORIES - else "" - ), - "monitored_directories": monitored_dirs, + "infinite_scroll": PromptManagerConfig.INFINITE_SCROLL, + "worker_threads": PromptManagerConfig.WORKER_THREADS, + "cpu_count": PromptManagerConfig.max_worker_threads(), + "gallery_root_paths": root_paths, + "gallery_root_path": root_paths[0] if root_paths else "", + "monitored_directories": [ + self._public_path(d) for d in self._monitored_directories() + ], }, } ) except Exception as e: + self.logger.error(f"Get settings error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to get settings: {str(e)}"}, + {"success": False, "error": "Failed to get settings"}, status=500, ) + def _monitored_directories(self): + """Directories the running image monitor watches, else the configured roots.""" + from ..config import GalleryConfig + + monitor = _current_image_monitor() + if monitor is not None: + return list(getattr(monitor, "monitored_directories", [])) + return list(GalleryConfig.MONITORING_DIRECTORIES) + async def save_settings(self, request): - """Save settings.""" + """Save settings: the whole payload is validated before anything changes.""" try: - from ..config import PromptManagerConfig, GalleryConfig + from ..config import PromptManagerConfig - data = await request.json() - restart_required = False - - # Update in-memory config - if "result_timeout" in data: - try: - PromptManagerConfig.RESULT_TIMEOUT = validate_result_timeout( - data["result_timeout"] - ) - except ValueError as ve: - return web.json_response( - {"success": False, "error": str(ve)}, status=400 - ) - if "webui_display_mode" in data: - PromptManagerConfig.WEBUI_DISPLAY_MODE = data["webui_display_mode"] - - # Blocked system directories (shared by both path handlers) - blocked = [ - "/etc", - "/usr", - "/bin", - "/sbin", - "/boot", - "/proc", - "/sys", - "/dev", - "/var/log", - "/root", - "C:\\Windows", - "C:\\Program Files", - ] - - # Handle gallery root paths (array β€” preferred) - if "gallery_root_paths" in data: - new_paths = data["gallery_root_paths"] - if not isinstance(new_paths, list): - return web.json_response( - { - "success": False, - "error": "gallery_root_paths must be a list", - }, - status=400, - ) - - validated_paths = [] - for path_str in new_paths: - path_str = path_str.strip() - if not path_str: - continue - resolved = Path(path_str).resolve() - if not resolved.is_dir(): - return web.json_response( - { - "success": False, - "error": f"Path does not exist or is not a directory: {path_str}", - }, - status=400, - ) - for b in blocked: - if str(resolved).startswith(b): - return web.json_response( - { - "success": False, - "error": f"Cannot use system directory: {path_str}", - }, - status=400, - ) - validated_paths.append(path_str) - - old_paths = list(GalleryConfig.MONITORING_DIRECTORIES) - if validated_paths != old_paths: - GalleryConfig.MONITORING_DIRECTORIES = validated_paths - self._cached_output_dir = None - self._gallery_cache = {} - restart_required = True - - elif "gallery_root_path" in data: - # Backward compat: single path string - new_path = data["gallery_root_path"].strip() - old_path = ( - GalleryConfig.MONITORING_DIRECTORIES[0] - if GalleryConfig.MONITORING_DIRECTORIES - else "" - ) - - if new_path != old_path: - if new_path: - resolved = Path(new_path).resolve() - if not resolved.is_dir(): - return web.json_response( - { - "success": False, - "error": f"Gallery path does not exist or is not a directory: {new_path}", - }, - status=400, - ) - for b in blocked: - if str(resolved).startswith(b): - return web.json_response( - { - "success": False, - "error": "Gallery path cannot point to a system directory", - }, - status=400, - ) - GalleryConfig.MONITORING_DIRECTORIES = [new_path] - else: - GalleryConfig.MONITORING_DIRECTORIES = [] - self._cached_output_dir = None - self._gallery_cache = {} - restart_required = True - - # Save to config file for persistence - config_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) - config_file = os.path.join(config_dir, "config.json") - - config_data = { - "web_ui": { - "result_timeout": PromptManagerConfig.RESULT_TIMEOUT, - "webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE, - }, - "gallery": { - "monitoring": {"directories": GalleryConfig.MONITORING_DIRECTORIES} - }, - } + data, error_response = await _read_json(request) + if error_response is not None: + return error_response try: - with open(config_file, "w") as f: - json.dump(config_data, f, indent=2) - self.logger.info(f"Settings saved to {config_file}") - except Exception as save_err: - self.logger.warning(f"Could not save config file: {save_err}") + updates = self._validate_settings(data) + except ValueError as ve: + return web.json_response( + {"success": False, "error": str(ve)}, status=400 + ) + + new_roots, error = self._parse_gallery_roots(data) + if error: + return web.json_response({"success": False, "error": error}, status=400) + + for name, value in updates.items(): + setattr(PromptManagerConfig, name, value) + restart_required = False + if new_roots is not None: + restart_required = self._apply_gallery_roots(new_roots) + + self._persist_settings() return web.json_response( { @@ -542,329 +784,240 @@ class AdminRoutesMixin: } ) except Exception as e: + self.logger.error(f"Save settings error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to save settings: {str(e)}"}, + {"success": False, "error": "Failed to save settings"}, status=500, ) + @staticmethod + def _validate_settings(data): + """Scalar settings of a payload as ``{CONFIG_ATTR: value}``. + + Raises ValueError with a user-facing message on the first bad value; + nothing is applied until every value passed. + """ + from ..config import PromptManagerConfig + + updates = {} + if "result_timeout" in data: + updates["RESULT_TIMEOUT"] = validate_result_timeout(data["result_timeout"]) + if "webui_display_mode" in data: + mode = data["webui_display_mode"] + if mode not in ("popup", "newtab"): + raise ValueError("webui_display_mode must be 'popup' or 'newtab'") + updates["WEBUI_DISPLAY_MODE"] = mode + if "infinite_scroll" in data: + if not isinstance(data["infinite_scroll"], bool): + raise ValueError("infinite_scroll must be true or false") + updates["INFINITE_SCROLL"] = data["infinite_scroll"] + if "worker_threads" in data: + updates["WORKER_THREADS"] = validate_worker_threads( + data["worker_threads"], PromptManagerConfig.max_worker_threads() + ) + return updates + + def _parse_gallery_roots(self, data): + """Extract and validate gallery roots from a settings payload. + + Accepts ``gallery_root_paths`` (list, preferred) or the legacy single + ``gallery_root_path`` string. Blank entries are dropped. + + Returns: + (roots, error): ``roots`` is the validated list, or ``None`` when + the payload carries no gallery root key; ``error`` is a message + when validation failed (and ``roots`` is then ``None``). + """ + if "gallery_root_paths" in data: + raw = data["gallery_root_paths"] + if not isinstance(raw, list): + return None, "gallery_root_paths must be a list" + elif "gallery_root_path" in data: + raw = [data["gallery_root_path"]] + else: + return None, None + + from ..config import GalleryConfig + + roots = [] + for entry in raw: + if not isinstance(entry, str): + return None, "Gallery root paths must be strings" + entry = entry.strip() + if not entry: + continue + ok, reason = GalleryConfig.validate_gallery_root(entry) + if not ok: + return ( + None, + f"Invalid gallery root '{os.path.basename(entry)}': {reason}", + ) + roots.append(GalleryConfig.resolve_gallery_root(entry)) + return roots, None + + def _apply_gallery_roots(self, roots): + """Install validated gallery roots; returns True when they changed.""" + from ..config import GalleryConfig + + if roots == list(GalleryConfig.MONITORING_DIRECTORIES): + return False + GalleryConfig.MONITORING_DIRECTORIES = roots + self._cached_output_dir = None + self._gallery_cache = {} + return True + + def _persist_settings(self): + """Merge the user-editable settings into the configured config.json. + + Only the keys this endpoint owns are rewritten; every other section + (integrations, database, performance, ...) and every other key in + ``web_ui``/``gallery`` is kept as it was. A config file that cannot + be parsed is replaced. + """ + from ..config import GalleryConfig, PromptManagerConfig, write_private_json + + config_file = PromptManagerConfig.get_config_path() + existing = self._read_config_file(config_file) + web_ui = dict(existing.get("web_ui") or {}) + web_ui["result_timeout"] = PromptManagerConfig.RESULT_TIMEOUT + web_ui["webui_display_mode"] = PromptManagerConfig.WEBUI_DISPLAY_MODE + web_ui["infinite_scroll"] = PromptManagerConfig.INFINITE_SCROLL + gallery = dict(existing.get("gallery") or {}) + monitoring = dict(gallery.get("monitoring") or {}) + monitoring["directories"] = list(GalleryConfig.MONITORING_DIRECTORIES) + gallery["monitoring"] = monitoring + performance = dict(existing.get("performance") or {}) + performance["worker_threads"] = PromptManagerConfig.WORKER_THREADS + config_data = { + **existing, + "web_ui": web_ui, + "gallery": gallery, + "performance": performance, + } + try: + write_private_json(config_file, config_data) + self.logger.info(f"Settings saved to {config_file}") + except OSError as save_err: + self.logger.warning(f"Could not save config file: {save_err}") + + def _read_config_file(self, config_file): + """Current contents of config.json as a dict, or {} when unusable.""" + try: + with open(config_file, "r", encoding="utf-8") as f: + data = json.load(f) + except (OSError, ValueError): + return {} + return data if isinstance(data, dict) else {} + async def run_diagnostics(self, request): """Run comprehensive system diagnostics and health checks.""" try: - results = {} + from ..config import GalleryConfig - # Check database - try: - db_path = self.db.model.db_path - if os.path.exists(db_path): - with sqlite3.connect(db_path) as conn: - conn.row_factory = sqlite3.Row - cursor = conn.execute("SELECT COUNT(*) as count FROM prompts") - prompt_count = cursor.fetchone()["count"] - - cursor = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='generated_images'" - ) - has_images_table = cursor.fetchone() is not None - - if has_images_table: - cursor = conn.execute( - "SELECT COUNT(*) as count FROM generated_images" - ) - image_count = cursor.fetchone()["count"] - else: - image_count = 0 - - results["database"] = { - "status": "ok", - "prompt_count": prompt_count, - "has_images_table": has_images_table, - "image_count": image_count, - } - else: - results["database"] = { - "status": "error", - "message": f"Database file not found: {db_path}", - } - except Exception as e: - results["database"] = { - "status": "error", - "message": f"Database error: {str(e)}", - } - - # Check dependencies - dependencies = {} - try: - import watchdog - - dependencies["watchdog"] = True - except ImportError: - dependencies["watchdog"] = False - - try: - from PIL import Image - - dependencies["PIL"] = True - except ImportError: - dependencies["PIL"] = False - - dependencies["sqlite3"] = True # Always available in Python - - results["dependencies"] = { - "status": "ok" if all(dependencies.values()) else "error", - "dependencies": dependencies, + results = { + "database": _diagnose_database(self.db.model.db_path), + "dependencies": _diagnose_dependencies(), + "comfyui_output": _diagnose_output_dirs( + GalleryConfig.MONITORING_DIRECTORIES + ), + "image_monitor": _diagnose_image_monitor(), } - - # Check output directories - output_dirs = [] - - # Check user-configured gallery directories from config.json - try: - from ..config import GalleryConfig - - for cfg_dir in GalleryConfig.MONITORING_DIRECTORIES: - abs_path = os.path.abspath(cfg_dir) - if os.path.exists(abs_path) and abs_path not in output_dirs: - output_dirs.append(abs_path) - except Exception as e: - self.logger.debug(f"Could not load gallery config: {e}") - - # Check ComfyUI's own output directory (respects --output-directory) - try: - import folder_paths - - comfyui_output = folder_paths.get_output_directory() - if comfyui_output and os.path.exists(comfyui_output): - abs_path = os.path.abspath(comfyui_output) - if abs_path not in output_dirs: - output_dirs.append(abs_path) - except ImportError: - self.logger.debug( - "folder_paths not available, skipping ComfyUI output dir" - ) - - # Fallback: check common relative paths - if not output_dirs: - for dir_path in ["output", "../output", "../../output"]: - abs_path = os.path.abspath(dir_path) - if os.path.exists(abs_path) and abs_path not in output_dirs: - output_dirs.append(abs_path) - - results["comfyui_output"] = { - "status": "ok" if output_dirs else "warning", - "output_dirs": output_dirs, - } - - # Check image monitor status - try: - from ...utils import image_monitor as im_mod - - monitor = im_mod._monitor_instance - if monitor is not None: - monitor_status = monitor.get_status() - results["image_monitor"] = { - "status": ( - "ok" if monitor_status.get("observer_alive") else "error" - ), - **monitor_status, - } - else: - results["image_monitor"] = { - "status": "error", - "message": "Image monitor not initialized", - } - except Exception as e: - results["image_monitor"] = { - "status": "error", - "message": f"Failed to get monitor status: {str(e)}", - } - return web.json_response({"success": True, "diagnostics": results}) except Exception as e: self.logger.error(f"Diagnostics error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": "Diagnostics failed"}, status=500 + ) async def test_image_link(self, request): """Test creating an image link.""" try: - data = await request.json() + data, error_response = await _read_json(request) + if error_response is not None: + return error_response prompt_id = data.get("prompt_id") - test_image_path = data.get("image_path", "/test/fake/image.png") - if not prompt_id: return web.json_response( {"success": False, "error": "prompt_id is required"}, status=400 ) - test_metadata = { - "file_info": { - "size": 1024000, - "dimensions": [512, 512], - "format": "PNG", - }, - "workflow": {"test": True}, - "prompt": {"test_prompt": "This is a test image"}, - } - - try: - image_id = await self._run_in_executor( - self.db.link_image_to_prompt, - prompt_id=str(prompt_id), - image_path=test_image_path, - metadata=test_metadata, - ) - - return web.json_response( - { - "success": True, - "result": { - "status": "ok", - "image_id": image_id, - "message": f"Test image linked successfully with ID {image_id}", - }, - } - ) - except Exception as e: - return web.json_response( - { - "success": False, - "result": { - "status": "error", - "message": f"Failed to create test link: {str(e)}", - }, - } - ) + # The client never chooses the stored path: a synthetic marker + # under the output directory stands in for a generated image. + output_dir = await self._run_in_executor(self._find_comfyui_output_dir) + image_path = self._test_link_marker_path(output_dir) + payload = await self._link_test_image(str(prompt_id), image_path) + return web.json_response(payload) except Exception as e: self.logger.error(f"Test link error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": "Test link failed"}, status=500 + ) + + TEST_LINK_MARKER = "prompt_manager_test_link.png" + + @classmethod + def _test_link_marker_path(cls, output_dir): + """Path recorded by the test-link diagnostic: a marker under ``output_dir`` + (relative, when no output directory is known).""" + if output_dir: + return os.path.join(str(output_dir), cls.TEST_LINK_MARKER) + return cls.TEST_LINK_MARKER + + async def _link_test_image(self, prompt_id, image_path): + """Link a synthetic image record to ``prompt_id``; returns the envelope.""" + test_metadata = { + "file_info": {"size": 1024000, "dimensions": [512, 512], "format": "PNG"}, + "workflow": {"test": True}, + "prompt": {"test_prompt": "This is a test image"}, + } + try: + image_id = await self._run_in_executor( + self.db.link_image_to_prompt, + prompt_id=prompt_id, + image_path=image_path, + metadata=test_metadata, + ) + except Exception as e: + return { + "success": False, + "result": { + "status": "error", + "message": f"Failed to create test link: {self._public_error(e)}", + }, + } + return { + "success": True, + "result": { + "status": "ok", + "image_id": image_id, + "message": f"Test image linked successfully with ID {image_id}", + }, + } + + DEFAULT_MAINTENANCE_OPERATIONS = ( + "cleanup_duplicates", + "vacuum", + "cleanup_orphaned_images", + ) async def run_maintenance(self, request): - """Perform comprehensive database maintenance and optimization.""" + """Perform database maintenance operations and report each outcome.""" try: - data = ( - await request.json() - if request.content_type == "application/json" - else {} - ) + data = {} + if request.content_type == "application/json": + data, error_response = await _read_json(request) + if error_response is not None: + return error_response operations = data.get( - "operations", - ["cleanup_duplicates", "vacuum", "cleanup_orphaned_images"], + "operations", list(self.DEFAULT_MAINTENANCE_OPERATIONS) ) - results = {} - - def _run_maintenance(): - if "cleanup_duplicates" in operations: - try: - duplicates_removed = self.db.cleanup_duplicates() - results["cleanup_duplicates"] = { - "success": True, - "removed_count": duplicates_removed, - "message": f"Removed {duplicates_removed} duplicate prompts", - } - except Exception as e: - results["cleanup_duplicates"] = { - "success": False, - "error": str(e), - "message": "Failed to cleanup duplicates", - } - - if "vacuum" in operations: - try: - self.db.model.vacuum_database() - results["vacuum"] = { - "success": True, - "message": "Database vacuum completed successfully", - } - except Exception as e: - results["vacuum"] = { - "success": False, - "error": str(e), - "message": "Failed to vacuum database", - } - - if "cleanup_orphaned_images" in operations: - try: - orphaned_removed = self.db.cleanup_missing_images() - results["cleanup_orphaned_images"] = { - "success": True, - "removed_count": orphaned_removed, - "message": f"Removed {orphaned_removed} orphaned image records", - } - except Exception as e: - results["cleanup_orphaned_images"] = { - "success": False, - "error": str(e), - "message": "Failed to cleanup orphaned images", - } - - if "check_hash_duplicates" in operations: - try: - hash_duplicates = self.db.check_hash_duplicates() - results["check_hash_duplicates"] = { - "success": True, - "duplicate_hashes": len(hash_duplicates), - "message": f"Found {len(hash_duplicates)} duplicate hash groups", - } - except Exception as e: - results["check_hash_duplicates"] = { - "success": False, - "error": str(e), - "message": "Failed to check hash duplicates", - } - - if "statistics" in operations: - try: - db_info = self.db.model.get_database_info() - results["statistics"] = { - "success": True, - "info": db_info, - "message": "Database statistics retrieved", - } - except Exception as e: - results["statistics"] = { - "success": False, - "error": str(e), - "message": "Failed to get database statistics", - } - - if "prune_orphaned_prompts" in operations: - try: - removed_count = self.db.prune_orphaned_prompts() - results["prune_orphaned_prompts"] = { - "success": True, - "removed_count": removed_count, - "message": f"Removed {removed_count} orphaned prompts (prompts with no linked images, excluding protected prompts)", - } - except Exception as e: - results["prune_orphaned_prompts"] = { - "success": False, - "error": str(e), - "message": "Failed to prune orphaned prompts", - } - - if "check_consistency" in operations: - try: - consistency_issues = self.db.check_consistency() - results["check_consistency"] = { - "success": True, - "issues_found": len(consistency_issues), - "issues": consistency_issues[:10], - "message": f"Found {len(consistency_issues)} consistency issues", - } - except Exception as e: - results["check_consistency"] = { - "success": False, - "error": str(e), - "message": "Failed to check database consistency", - } - - await self._run_in_executor(_run_maintenance) - - all_successful = all( - result.get("success", False) for result in results.values() + results = await self._run_in_executor( + self._run_maintenance_operations, operations ) + all_successful = all(r.get("success", False) for r in results.values()) return web.json_response( { @@ -872,35 +1025,120 @@ class AdminRoutesMixin: "operations_completed": len(results), "all_successful": all_successful, "results": results, - "message": f"Maintenance completed: {len(results)} operations processed", + "message": ( + f"Maintenance completed: {len(results)} operations processed" + ), } ) except Exception as e: self.logger.error(f"Maintenance error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Maintenance failed: {str(e)}"}, status=500 + {"success": False, "error": "Maintenance failed"}, status=500 ) + def _run_maintenance_operations(self, operations): + """Run each known operation in order; unknown names are ignored (blocking).""" + results = {} + for name in operations: + runner = self._maintenance_runners().get(name) + if runner is None: + continue + try: + results[name] = {"success": True, **runner()} + except Exception as e: + results[name] = { + "success": False, + "error": self._public_error(e), + "message": f"Failed to run {name.replace('_', ' ')}", + } + return results + + def _maintenance_runners(self): + """Map of operation name to a callable returning that operation's result.""" + db = self.db + + def count_result(count, noun): + return {"removed_count": count, "message": f"Removed {count} {noun}"} + + def vacuum(): + db.model.vacuum_database() + return {"message": "Database vacuum completed successfully"} + + return { + "cleanup_duplicates": lambda: count_result( + db.cleanup_duplicates(), "duplicate prompts" + ), + "vacuum": vacuum, + "cleanup_orphaned_images": lambda: count_result( + db.cleanup_missing_images(), "orphaned image records" + ), + "check_hash_duplicates": lambda: self._hash_duplicates_result( + db.check_hash_duplicates() + ), + "statistics": lambda: { + "info": self._public_database_info(db.model.get_database_info()), + "message": "Database statistics retrieved", + }, + "prune_orphaned_prompts": lambda: count_result( + db.prune_orphaned_prompts(), + "orphaned prompts (prompts with no linked images, " + "excluding protected prompts)", + ), + "check_consistency": lambda: self._consistency_result( + db.check_consistency() + ), + } + + def _public_database_info(self, info): + """``get_database_info()`` with the absolute database_path made public.""" + info = dict(info or {}) + if info.get("database_path"): + info["database_path"] = self._public_path(info["database_path"]) + return info + + @staticmethod + def _hash_duplicates_result(groups): + return { + "duplicate_hashes": len(groups), + "message": f"Found {len(groups)} duplicate hash groups", + } + + @staticmethod + def _consistency_result(issues): + return { + "issues_found": len(issues), + "issues": issues[:10], + "message": f"Found {len(issues)} consistency issues", + } + async def backup_database(self, request): - """Backup the entire prompts.db database file.""" + """Download a consistent copy of the prompts database (WAL included).""" try: - db_path = self.db.model.db_path + model = self.db.model + db_path = model.db_path if not os.path.exists(db_path): return web.json_response( {"success": False, "error": "Database file not found"}, status=404 ) - with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as temp_file: - temp_path = temp_file.name - - shutil.copy2(db_path, temp_path) - - with open(temp_path, "rb") as f: - file_data = f.read() - - os.unlink(temp_path) + fd, temp_path = tempfile.mkstemp(suffix=".db") + os.close(fd) + try: + ok = await self._run_in_executor(model.backup_database, temp_path) + if not ok: + return web.json_response( + {"success": False, "error": "Failed to create database backup"}, + status=500, + ) + # Read once (the file is deleted right after) rather than + # streaming so the temp file never outlives the request. + with open(temp_path, "rb") as fh: + file_data = fh.read() + finally: + if os.path.exists(temp_path): + os.unlink(temp_path) timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"prompts_backup_{timestamp}.db" @@ -917,12 +1155,22 @@ class AdminRoutesMixin: except Exception as e: self.logger.error(f"Backup error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to backup database: {str(e)}"}, + { + "success": False, + "error": f"Failed to backup database: {self._public_error(e)}", + }, status=500, ) async def restore_database(self, request): - """Restore the prompts.db database from uploaded file.""" + """Restore the prompts database from an uploaded SQLite file. + + The upload is streamed to a temp file (capped at restore_max_bytes), + verified with PromptModel.verify_database_file and only then copied + over the live database, which is backed up first. + """ + max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024) + temp_path = None try: reader = await request.multipart() field = await reader.next() @@ -931,278 +1179,134 @@ class AdminRoutesMixin: return web.json_response( { "success": False, - "error": "No database file uploaded. Expected field name: database_file", + "error": ( + "No database file uploaded. " + "Expected field name: database_file" + ), }, status=400, ) - MAX_RESTORE_SIZE = 100 * 1024 * 1024 # 100MB - file_data = await field.read() + fd, temp_path = tempfile.mkstemp(suffix=".db") + os.close(fd) + size = 0 + with open(temp_path, "wb") as out: + while True: + chunk = await field.read_chunk() + if not chunk: + break + size += len(chunk) + if size > max_bytes: + return web.json_response( + { + "success": False, + "error": ( + "File too large. Maximum size is " + f"{max_bytes // (1024 * 1024)}MB" + ), + }, + status=400, + ) + out.write(chunk) - if not file_data: + if size == 0: return web.json_response( {"success": False, "error": "Uploaded file is empty"}, status=400 ) - if len(file_data) > MAX_RESTORE_SIZE: + model = self.db.model + ok, reason = await self._run_in_executor( + model.verify_database_file, temp_path + ) + if ok: + reason = await self._run_in_executor(_inspect_restore_upload, temp_path) + ok = reason is None + if not ok: return web.json_response( - { - "success": False, - "error": f"File too large. Maximum size is {MAX_RESTORE_SIZE // (1024*1024)}MB", - }, - status=400, + {"success": False, "error": reason}, status=400 ) - with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as temp_file: - temp_path = temp_file.name - temp_file.write(file_data) + backup_path = await self._run_in_executor( + model.restore_from_file, temp_path + ) + removed = await self._run_in_executor(_prune_safety_backups, model.db_path) + if removed: + self.logger.info(f"Removed {len(removed)} old safety backups") + previous = self.db + self.db = PromptDatabase(model.db_path) + # Its connections would otherwise linger (a file lock on Windows) + previous.close_all() + prompt_count = self.db.model.get_database_info().get("total_prompts", 0) - try: - with sqlite3.connect(temp_path) as conn: - conn.row_factory = sqlite3.Row - - cursor = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'" - ) - if not cursor.fetchone(): - raise ValueError("Database does not contain a 'prompts' table") - - cursor = conn.execute("PRAGMA table_info(prompts)") - columns = [row["name"] for row in cursor.fetchall()] - required_columns = ["id", "text", "created_at"] - - for col in required_columns: - if col not in columns: - raise ValueError(f"Database missing required column: {col}") - - cursor = conn.execute("SELECT COUNT(*) as count FROM prompts") - prompt_count = cursor.fetchone()["count"] - - db_path = self.db.model.db_path - backup_path = f"{db_path}.backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" - - if os.path.exists(db_path): - shutil.copy2(db_path, backup_path) - self.logger.info(f"Current database backed up to: {backup_path}") - - shutil.copy2(temp_path, db_path) - - # Reinitialize the database connection - try: - from ...database.operations import PromptDatabase - except ImportError: - import sys - - sys.path.insert( - 0, - os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ), - ) - from database.operations import PromptDatabase - self.db = PromptDatabase() - - return web.json_response( - { - "success": True, - "message": f"Database restored successfully. Found {prompt_count} prompts.", - "prompt_count": prompt_count, - "backup_created": ( - backup_path if os.path.exists(db_path) else None - ), - } - ) - - except sqlite3.Error as e: - return web.json_response( - {"success": False, "error": f"Invalid SQLite database: {str(e)}"}, - status=400, - ) - except ValueError as e: - return web.json_response( - {"success": False, "error": str(e)}, status=400 - ) - finally: - if os.path.exists(temp_path): - os.unlink(temp_path) + return web.json_response( + { + "success": True, + "message": ( + "Database restored successfully. " + f"Found {prompt_count} prompts." + ), + "prompt_count": prompt_count, + "backup_created": ( + self._public_path(backup_path) if backup_path else None + ), + } + ) except Exception as e: self.logger.error(f"Restore error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to restore database: {str(e)}"}, + { + "success": False, + "error": f"Failed to restore database: {self._public_error(e)}", + }, status=500, ) + finally: + if temp_path: + for suffix in ("", "-wal", "-shm", "-journal"): + if os.path.exists(temp_path + suffix): + os.unlink(temp_path + suffix) async def scan_images(self, request): - """Scan ComfyUI output images for prompt metadata and add them to the database.""" + """Scan ComfyUI output images for prompt metadata and add them to the database. - BATCH_SIZE = 50 # Files per executor call for metadata extraction + The scan runs as a background task that outlives this request, so a + closed tab never leaves a half-finished scan. Progress streams as + server-sent events; while a scan is running, a new request attaches + to it instead of starting another one. + """ + job = self._scan_job + if job is None or job.finished: + job = self._start_scan_job() + return await self._stream_scan_job(request, job) - def _collect_media_files(output_dirs): - """Collect all media files from output directories (blocking I/O).""" - image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] - video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] - media_extensions = image_extensions + video_extensions - all_files = [] - seen = set() - for output_dir in output_dirs: - for ext in media_extensions: - for pattern in [f"*{ext}", f"*{ext.upper()}"]: - for f in output_dir.rglob(pattern): - if "thumbnails" not in f.parts: - norm = str(f).lower() - if norm not in seen: - seen.add(norm) - all_files.append(f) - return all_files + async def scan_status(self, request): + """Whether an output scan is running, and its latest progress event.""" + job = self._scan_job + return web.json_response( + { + "success": True, + "running": bool(job and job.running), + "last_event": job.last_event if job else None, + } + ) - def _extract_batch_metadata(file_batch): - """Extract metadata from a batch of files in one executor call.""" - results = [] - for f in file_batch: - try: - meta = self._extract_comfyui_metadata(str(f)) - results.append((f, meta)) - except Exception: - results.append((f, {})) - return results + def _start_scan_job(self): + """Create the job record synchronously, then run the scan as a task.""" + job = ScanJob() + self._scan_job = job + job.task = asyncio.ensure_future(self._run_scan_job(job)) + return job - async def stream_response(): - try: - self.logger.info("Starting image scan operation") - - output_dirs = self._get_all_output_dirs() - if not output_dirs: - self.logger.error("No output directories found") - yield f"data: {json.dumps({'type': 'error', 'message': 'No output directories found. Configure scan directories in Settings.'})}\n\n" - return - - dir_names = [str(d) for d in output_dirs] - yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': f'Scanning {len(output_dirs)} directory(ies) for media files...', 'processed': 0, 'found': 0})}\n\n" - - media_files = await self._run_in_executor( - _collect_media_files, output_dirs - ) - total_files = len(media_files) - - if total_files == 0: - yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'found': 0, 'added': 0, 'linked': 0, 'directories': dir_names})}\n\n" - return - - yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': f'Found {total_files} media files to process...', 'processed': 0, 'found': 0})}\n\n" - - processed_count = 0 - found_count = 0 - added_count = 0 - linked_count = 0 - - try: - from ...utils.hashing import generate_prompt_hash - except ImportError: - import sys - - current_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) - sys.path.insert(0, current_dir) - from utils.hashing import generate_prompt_hash - - # Process files in batches for performance - for batch_start in range(0, total_files, BATCH_SIZE): - batch = media_files[batch_start : batch_start + BATCH_SIZE] - - # Extract metadata for entire batch in one executor call - batch_results = await self._run_in_executor( - _extract_batch_metadata, batch - ) - - for media_file, metadata in batch_results: - try: - processed_count += 1 - - if not metadata: - continue - - parsed_data = self._parse_comfyui_prompt(metadata) - - if not ( - parsed_data.get("prompt") - or parsed_data.get("parameters") - ): - continue - - found_count += 1 - prompt_text = self._extract_readable_prompt(parsed_data) - - if prompt_text and not isinstance(prompt_text, str): - prompt_text = str(prompt_text) - - if not (prompt_text and prompt_text.strip()): - continue - - prompt_hash = generate_prompt_hash(prompt_text.strip()) - existing = await self._run_in_executor( - self.db.get_prompt_by_hash, prompt_hash - ) - - if existing: - try: - await self._run_in_executor( - self.db.link_image_to_prompt, - existing["id"], - str(media_file), - ) - linked_count += 1 - except Exception as e: - self.logger.error( - f"Failed to link {media_file.name} to existing prompt: {e}" - ) - else: - prompt_id = await self._run_in_executor( - self.db.save_prompt, - prompt_text.strip(), - "scanned", - ["auto-scanned"], - None, - f"Auto-scanned from {media_file.name}", - prompt_hash, - ) - if prompt_id: - added_count += 1 - try: - await self._run_in_executor( - self.db.link_image_to_prompt, - prompt_id, - str(media_file), - ) - except Exception as e: - self.logger.error( - f"Failed to link {media_file.name} to new prompt: {e}" - ) - - except Exception as e: - self.logger.error( - f"Error processing {media_file.name}: {e}" - ) - continue - - # Progress update after each batch - progress = int( - min(batch_start + len(batch), total_files) / total_files * 100 - ) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Processing file {min(batch_start + len(batch), total_files)}/{total_files}...', 'processed': processed_count, 'found': found_count})}\n\n" - await asyncio.sleep(0) - - self.logger.info( - f"Scan completed: processed={processed_count}, found={found_count}, " - f"new_prompts_added={added_count}, images_linked_to_existing={linked_count}" - ) - yield f"data: {json.dumps({'type': 'complete', 'processed': processed_count, 'found': found_count, 'added': added_count, 'linked': linked_count, 'directories': dir_names})}\n\n" - - except Exception as e: - self.logger.exception("Scan error") - yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\n\n" + async def _run_scan_job(self, job): + try: + async for payload in self._scan_images_events(): + job.publish(payload) + finally: + job.finish() + async def _stream_scan_job(self, request, job): + """Stream ``job`` as SSE; stop quietly if the client goes away.""" response = web.StreamResponse( status=200, reason="OK", @@ -1212,11 +1316,175 @@ class AdminRoutesMixin: "Connection": "keep-alive", }, ) - await response.prepare(request) - - async for chunk in stream_response(): - await response.write(chunk.encode("utf-8")) - - await response.write_eof() + try: + async with contextlib.aclosing(job.events()) as events: + async for payload in events: + await response.write(_sse(payload).encode("utf-8")) + await response.write_eof() + except Exception as e: # the socket closed under us; the scan goes on + self.logger.info( + f"Scan progress client disconnected ({e}); " + "the scan continues in the background" + ) return response + + async def _scan_images_events(self): + """Yield progress payloads while scanning every configured output directory.""" + try: + self.logger.info("Starting image scan operation") + output_dirs = self._get_all_output_dirs() + if not output_dirs: + self.logger.error("No output directories found") + yield { + "type": "error", + "message": "No output directories found. " + "Configure scan directories in Settings.", + } + return + + dir_names = [self._public_path(d) for d in output_dirs] + yield self._scan_progress( + 0, f"Scanning {len(output_dirs)} directory(ies) for media files..." + ) + + media_files = await self._run_in_executor( + self._collect_output_media_sync, output_dirs + ) + counts = {"processed": 0, "found": 0, "added": 0, "linked": 0} + if media_files: + yield self._scan_progress( + 5, f"Found {len(media_files)} media files to process...", counts + ) + async for event in self._scan_batches(media_files, counts): + yield event + + self.logger.info( + f"Scan completed: processed={counts['processed']}, " + f"found={counts['found']}, new_prompts_added={counts['added']}, " + f"images_linked_to_existing={counts['linked']}" + ) + yield {"type": "complete", **counts, "directories": dir_names} + + except Exception: + self.logger.exception("Scan error") + yield { + "type": "error", + "message": "An internal error occurred. Check server logs for details.", + } + + async def _scan_batches(self, media_files, counts): + """Process ``media_files`` in executor batches, updating ``counts`` in place. + + Yields one progress event per batch. + """ + total = len(media_files) + for batch_start in range(0, total, SCAN_BATCH_SIZE): + batch = media_files[batch_start : batch_start + SCAN_BATCH_SIZE] + batch_results = await self._run_in_executor( + self._extract_batch_metadata_sync, batch + ) + batch_counts = await self._run_in_executor( + self._ingest_scan_batch_sync, batch_results + ) + for key in counts: + counts[key] += batch_counts[key] + + done = min(batch_start + len(batch), total) + yield self._scan_progress( + int(done / total * 100), f"Processing file {done}/{total}...", counts + ) + await asyncio.sleep(0) + + @staticmethod + def _scan_progress(progress, status, counts=None): + """Progress event payload for the output scan.""" + counts = counts or {} + return { + "type": "progress", + "progress": progress, + "status": status, + "processed": counts.get("processed", 0), + "found": counts.get("found", 0), + } + + def _collect_output_media_sync(self, output_dirs): + """Collect all media files from the output directories (blocking I/O).""" + return _collect_media_files(output_dirs, IMAGE_EXTENSIONS + VIDEO_EXTENSIONS) + + def _extract_batch_metadata_sync(self, file_batch): + """Extract metadata from a batch of files in one executor call. + + Files are read on ``PromptManagerConfig.WORKER_THREADS`` threads: + Pillow decodes the whole image to reach trailing text chunks and + releases the GIL while doing so. + """ + from ..config import PromptManagerConfig + + def extract(media_file): + try: + return (media_file, self._extract_comfyui_metadata(str(media_file))) + except Exception: + return (media_file, {}) + + return map_parallel(extract, file_batch, PromptManagerConfig.WORKER_THREADS) + + def _ingest_scan_batch_sync(self, batch_results): + """Store the prompts found in a metadata batch (blocking DB work).""" + counts = {"processed": 0, "found": 0, "added": 0, "linked": 0} + for media_file, metadata in batch_results: + counts["processed"] += 1 + try: + outcome = self._ingest_scanned_file(media_file, metadata) + except Exception as e: + self.logger.error(f"Error processing {media_file.name}: {e}") + continue + if outcome is None: + continue + counts["found"] += 1 + if outcome in ("added", "linked"): + counts[outcome] += 1 + return counts + + def _ingest_scanned_file(self, media_file, metadata): + """Save or link one scanned file; returns 'added', 'linked', 'found' or None.""" + if not metadata: + return None + parsed = self._parse_comfyui_prompt(metadata) + if not (parsed.get("prompt") or parsed.get("parameters")): + return None + + prompt_text = self._extract_readable_prompt(parsed) + if prompt_text is not None and not isinstance(prompt_text, str): + prompt_text = str(prompt_text) + if not (prompt_text and prompt_text.strip()): + return "found" + prompt_text = prompt_text.strip() + + prompt_hash = generate_prompt_hash(prompt_text) + existing = self.db.get_prompt_by_hash(prompt_hash) + if existing: + try: + self.db.link_image_to_prompt(existing["id"], str(media_file)) + except Exception as e: + self.logger.error( + f"Failed to link {media_file.name} to existing prompt: {e}" + ) + return "found" + return "linked" + + prompt_id = self.db.save_prompt( + prompt_text, + "scanned", + ["auto-scanned"], + None, + f"Auto-scanned from {media_file.name}", + prompt_hash, + ) + if not prompt_id: + return "found" + try: + self.db.link_image_to_prompt(prompt_id, str(media_file)) + except Exception as e: + self.logger.error(f"Failed to link {media_file.name} to new prompt: {e}") + return "added" diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 4e0a1cf..dd3019c 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -2,10 +2,53 @@ import asyncio import json +import os +import threading from pathlib import Path +from typing import Optional +from urllib.parse import quote + +from .prompts import thumbnail_url_for from aiohttp import web +_INTERNAL_ERROR = "An internal error occurred. Check server logs for details." +_BUSY_ERROR = "Auto-tagging is already running" + +# Only these may be handed to the tagging engine or listed by the scanner. +_MEDIA_EXTENSIONS = frozenset( + {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"} +) +# Bounds for scan_output_dir: a runaway output tree must not pin the server. +_SCAN_MAX_FILES = 50_000 +_SCAN_MAX_DEPTH = 12 + +# One batch run at a time; a second autotag/start gets 409 while this is held. +_autotag_batch_lock = threading.Lock() + + +def _file_exists(path) -> bool: + try: + return Path(str(path)).exists() + except (OSError, ValueError): + return False + + +def path_is_within(candidate, root) -> bool: + """True when ``candidate`` is strictly inside ``root`` on the real filesystem. + + Both sides go through ``os.path.realpath`` (symlinks, ``..``) and + ``os.path.normcase`` (Windows case/separator folding) before the + ``Path.is_relative_to`` check, so ``/out2`` is never treated as inside + ``/out`` and a symlink inside ``root`` pointing elsewhere is rejected. + """ + try: + child = Path(os.path.normcase(os.path.realpath(str(candidate)))) + parent = Path(os.path.normcase(os.path.realpath(str(root)))) + except (OSError, RuntimeError, ValueError): + return False + return child != parent and child.is_relative_to(parent) + class AutotagRoutesMixin: """Mixin providing auto-tagging API endpoints.""" @@ -15,11 +58,11 @@ class AutotagRoutesMixin: async def get_autotag_models_route(request): return await self.get_autotag_models(request) - @routes.get("/prompt_manager/autotag/download/{model_type}") + @routes.post("/prompt_manager/autotag/download/{model_type}") async def download_autotag_model_route(request): return await self.download_autotag_model(request) - @routes.get("/prompt_manager/autotag/start") + @routes.post("/prompt_manager/autotag/start") async def start_autotag_route(request): return await self.start_autotag(request) @@ -61,253 +104,17 @@ class AutotagRoutesMixin: except Exception as e: self.logger.error(f"Get autotag models error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) - - async def download_autotag_model(self, request): - """Download an AutoTag model with streaming progress.""" - model_type = request.match_info.get("model_type") - - async def stream_response(): - try: - from ..autotag import get_autotag_service - - service = get_autotag_service() - - if model_type not in service.models_config: - yield f"data: {json.dumps({'type': 'error', 'message': f'Invalid model type: {model_type}'})}\n\n" - return - - yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Starting download...'})}\n\n" - - progress_data = {"last_progress": 0} - - def progress_callback(status: str, progress: float): - progress_data["last_progress"] = progress - - loop = asyncio.get_event_loop() - success = await loop.run_in_executor( - None, lambda: service.download_model(model_type, progress_callback) - ) - - if success: - yield f"data: {json.dumps({'type': 'complete', 'progress': 100, 'status': 'Download complete'})}\n\n" - else: - yield f"data: {json.dumps({'type': 'error', 'message': 'Download failed'})}\n\n" - - except Exception as e: - self.logger.exception("Download model error") - yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\n\n" - - response = web.StreamResponse( - status=200, - reason="OK", - headers={ - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Connection": "keep-alive", - }, - ) - - await response.prepare(request) - - async for chunk in stream_response(): - await response.write(chunk.encode("utf-8")) - - await response.write_eof() - return response - - async def start_autotag(self, request): - """Start batch auto-tagging with streaming progress.""" - model_type = request.query.get("model_type", "gguf") - custom_prompt = request.query.get("prompt", "") - skip_tagged = request.query.get("skip_tagged", "true").lower() == "true" - keep_in_memory = request.query.get("keep_in_memory", "true").lower() == "true" - use_gpu = True - - # WD14 threshold params - general_threshold = request.query.get("general_threshold") - character_threshold = request.query.get("character_threshold") - try: - if general_threshold is not None: - general_threshold = float(general_threshold) - if character_threshold is not None: - character_threshold = float(character_threshold) - except (ValueError, TypeError): return web.json_response( - { - "success": False, - "error": "general_threshold and character_threshold must be numeric", - }, - status=400, + {"success": False, "error": self._public_error(e)}, status=500 ) - async def stream_response(): - try: - from ..autotag import get_autotag_service - import time as _time - - service = get_autotag_service() - - status = service.get_models_status() - if not status.get(model_type, {}).get("downloaded"): - yield f"data: {json.dumps({'type': 'error', 'message': f'Model {model_type} not downloaded'})}\n\n" - return - - yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Loading model...'})}\n\n" - - loop = asyncio.get_event_loop() - try: - await loop.run_in_executor( - None, lambda: service.load_model(model_type, use_gpu) - ) - except Exception as e: - self.logger.exception("Failed to load model") - yield f"data: {json.dumps({'type': 'error', 'message': 'Failed to load model. Check server logs for details.'})}\n\n" - return - - if custom_prompt: - service.custom_prompt = custom_prompt - - yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': 'Model loaded. Fetching all images from database...'})}\n\n" - - images = await self._run_in_executor(self.db.get_all_images) - - total_files = len(images) - if total_files == 0: - yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'tagged': 0, 'skipped': 0, 'status': 'No images with linked prompts found'})}\n\n" - service.unload_model() - return - - yield f"data: {json.dumps({'type': 'progress', 'progress': 10, 'status': f'Found {total_files} images. Processing...'})}\n\n" - - processed = 0 - tagged = 0 - skipped = 0 - errors = 0 - tagged_prompt_ids = set() - - last_update_time = _time.monotonic() - - for i, image_data in enumerate(images): - image_path = image_data.get("image_path") - prompt_id = image_data.get("prompt_id") - - if not image_path or not prompt_id: - skipped += 1 - continue - - if prompt_id in tagged_prompt_ids: - skipped += 1 - now = _time.monotonic() - if (now - last_update_time) >= 0.5 or i == total_files - 1: - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (prompt already processed)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = now - continue - - if not Path(image_path).exists(): - skipped += 1 - now = _time.monotonic() - if (now - last_update_time) >= 0.5 or i == total_files - 1: - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (file missing)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = now - continue - - if skip_tagged: - prompt_tags = image_data.get("prompt_tags", []) - if isinstance(prompt_tags, str): - prompt_tags = [ - t.strip() for t in prompt_tags.split(",") if t.strip() - ] - real_tags = [ - t - for t in prompt_tags - if t != "auto-scanned" - and not t.startswith("prepend:") - and not t.startswith("append:") - ] - if real_tags: - tagged_prompt_ids.add(prompt_id) - skipped += 1 - now = _time.monotonic() - if (now - last_update_time) >= 0.5 or i == total_files - 1: - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (already tagged)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = now - continue - - try: - tags = await loop.run_in_executor( - None, - lambda p=str(image_path): service.generate_tags( - p, - general_threshold=general_threshold, - character_threshold=character_threshold, - ), - ) - - processed += 1 - - if tags: - existing_prompt = await self._run_in_executor( - self.db.get_prompt_by_id, prompt_id - ) - if existing_prompt: - existing_tags = existing_prompt.get("tags", []) - if isinstance(existing_tags, str): - existing_tags = [ - t.strip() - for t in existing_tags.split(",") - if t.strip() - ] - - new_tags = [t for t in tags if t not in existing_tags] - if new_tags: - all_tags = existing_tags + new_tags - await self._run_in_executor( - self.db.update_prompt_metadata, - prompt_id, - tags=all_tags, - ) - tagged += 1 - else: - skipped += 1 - else: - skipped += 1 - else: - skipped += 1 - - tagged_prompt_ids.add(prompt_id) - - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Processing {i+1}/{total_files}...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = _time.monotonic() - - except Exception as img_err: - self.logger.error(f"Error processing {image_path}: {img_err}") - errors += 1 - processed += 1 - - if not keep_in_memory: - service.unload_model() - model_status = "Model unloaded" - else: - model_status = "Model kept in memory" - - yield f"data: {json.dumps({'type': 'complete', 'progress': 100, 'processed': processed, 'tagged': tagged, 'skipped': skipped, 'errors': errors, 'status': 'Complete', 'model_status': model_status, 'model_loaded': keep_in_memory})}\n\n" - - except Exception as e: - self.logger.error(f"AutoTag error: {e}") - import traceback - - self.logger.exception("AutoTag error") - yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\n\n" + async def _stream_sse(self, request, events): + """Write an async iterator of dict events as server-sent events. + A client disconnect raises ConnectionResetError from ``write``; the + event generator is then closed so its cleanup runs before re-raising + (aiohttp logs premature disconnects at debug level). + """ response = web.StreamResponse( status=200, reason="OK", @@ -317,20 +124,393 @@ class AutotagRoutesMixin: "Connection": "keep-alive", }, ) - await response.prepare(request) - - async for chunk in stream_response(): - await response.write(chunk.encode("utf-8")) - + try: + async for event in events: + await response.write(f"data: {json.dumps(event)}\n\n".encode("utf-8")) + except (ConnectionResetError, asyncio.CancelledError): + self.logger.info("Client disconnected; stopping autotag stream") + await events.aclose() + raise await response.write_eof() return response + async def download_autotag_model(self, request): + """Download an AutoTag model with streaming progress (POST only).""" + model_type = request.match_info.get("model_type") + + async def events(): + try: + from ..autotag import get_autotag_service + + service = get_autotag_service() + if model_type not in service.models_config: + yield { + "type": "error", + "message": f"Invalid model type: {model_type}", + } + return + + yield { + "type": "progress", + "progress": 0, + "status": "Starting download...", + } + + def progress_callback(status: str, progress: float): + pass + + success = await self._run_in_executor( + service.download_model, model_type, progress_callback + ) + if success: + yield { + "type": "complete", + "progress": 100, + "status": "Download complete", + } + else: + yield {"type": "error", "message": "Download failed"} + except Exception: + self.logger.exception("Download model error") + yield { + "type": "error", + "message": _INTERNAL_ERROR, + } + + return await self._stream_sse(request, events()) + + async def _read_start_params(self, request) -> Optional[dict]: + """Merge query-string and optional JSON-body parameters for autotag/start. + + Returns None when a body is present but is not a JSON object. + """ + params = dict(request.query) + if request.can_read_body: + try: + body = await request.json() + except (json.JSONDecodeError, UnicodeDecodeError): + return None + if not isinstance(body, dict): + return None + params.update(body) + return params + + @staticmethod + def _real_tags(tags) -> list: + """Normalise a tag field to a list without bookkeeping markers.""" + if isinstance(tags, str): + tags = [t.strip() for t in tags.split(",") if t.strip()] + return [ + t + for t in (tags or []) + if t != "auto-scanned" + and not t.startswith("prepend:") + and not t.startswith("append:") + ] + + async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool: + """Merge ``tags`` into the prompt; True when something new was added.""" + existing_prompt = await self._run_in_executor( + self.db.get_prompt_by_id, prompt_id + ) + if not existing_prompt: + return False + existing_tags = existing_prompt.get("tags", []) + if isinstance(existing_tags, str): + existing_tags = [t.strip() for t in existing_tags.split(",") if t.strip()] + new_tags = [t for t in tags if t not in existing_tags] + if not new_tags: + return False + await self._run_in_executor( + self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags + ) + return True + + async def _autotag_batch_events(self, service, opts: dict): + """Yield SSE progress dicts while tagging every image linked to a prompt.""" + import time as _time + + model_type = opts["model_type"] + keep_in_memory = opts["keep_in_memory"] + status = service.get_models_status() + if not status.get(model_type, {}).get("downloaded"): + yield {"type": "error", "message": f"Model {model_type} not downloaded"} + return + + yield {"type": "progress", "progress": 0, "status": "Loading model..."} + try: + await self._run_in_executor(service.load_model, model_type, True) + except Exception: + self.logger.exception("Failed to load model") + yield { + "type": "error", + "message": "Failed to load model. Check server logs for details.", + } + return + + if opts["custom_prompt"]: + service.custom_prompt = opts["custom_prompt"] + + yield { + "type": "progress", + "progress": 5, + "status": "Model loaded. Fetching all images from database...", + } + images = await self._run_in_executor(self.db.get_all_images) + total = len(images) + if total == 0: + yield { + "type": "complete", + "processed": 0, + "tagged": 0, + "skipped": 0, + "errors": 0, + "skipped_outside_roots": 0, + "status": "No images with linked prompts found", + } + service.unload_model() + return + + yield { + "type": "progress", + "progress": 10, + "status": f"Found {total} images. Processing...", + } + + counts = { + "processed": 0, + "tagged": 0, + "skipped": 0, + "errors": 0, + "skipped_outside_roots": 0, + } + tagged_prompt_ids = set() + last_update = _time.monotonic() + finished = False + + def progress_event(i, message): + return { + "type": "progress", + "progress": 10 + int((i + 1) / total * 85), + "status": message, + "processed": counts["processed"], + "tagged": counts["tagged"], + "skipped": counts["skipped"], + "skipped_outside_roots": counts["skipped_outside_roots"], + } + + try: + for i, image_data in enumerate(images): + image_path = image_data.get("image_path") + prompt_id = image_data.get("prompt_id") + skip_reason = None + allowed_path = None + if not image_path or not prompt_id: + skip_reason = "unlinked" + elif prompt_id in tagged_prompt_ids: + skip_reason = "prompt already processed" + else: + # Rows can point anywhere (restored DB, hand edits): only + # media files inside the served output roots are tagged. + allowed_path = self._resolve_allowed_image_path(str(image_path)) + if allowed_path is None: + if _file_exists(image_path): + skip_reason = "outside allowed directories" + counts["skipped_outside_roots"] += 1 + else: + skip_reason = "file missing" + elif opts["skip_tagged"] and self._real_tags( + image_data.get("prompt_tags", []) + ): + tagged_prompt_ids.add(prompt_id) + skip_reason = "already tagged" + + if skip_reason is not None: + counts["skipped"] += 1 + now = _time.monotonic() + if (now - last_update) >= 0.5 or i == total - 1: + yield progress_event( + i, f"Skipping {i + 1}/{total} ({skip_reason})..." + ) + await asyncio.sleep(0.01) + last_update = now + continue + + try: + tags = await self._run_in_executor( + service.generate_tags, + allowed_path, + general_threshold=opts["general_threshold"], + character_threshold=opts["character_threshold"], + ) + counts["processed"] += 1 + if tags and await self._apply_tags_to_prompt(prompt_id, tags): + counts["tagged"] += 1 + else: + counts["skipped"] += 1 + tagged_prompt_ids.add(prompt_id) + except Exception as img_err: + self.logger.error( + f"Error processing {os.path.basename(str(image_path))}: " + f"{img_err}" + ) + counts["errors"] += 1 + counts["processed"] += 1 + + yield progress_event(i, f"Processing {i + 1}/{total}...") + await asyncio.sleep(0.01) + last_update = _time.monotonic() + + finished = True + if not keep_in_memory: + service.unload_model() + yield { + "type": "complete", + "progress": 100, + **counts, + "status": "Complete", + "model_status": ( + "Model kept in memory" if keep_in_memory else "Model unloaded" + ), + "model_loaded": keep_in_memory, + } + finally: + if not finished and not keep_in_memory: + service.unload_model() + + async def start_autotag(self, request): + """Start batch auto-tagging with streaming progress (POST only).""" + params = await self._read_start_params(request) + if params is None: + return self._bad_request("Request body must be a JSON object") + + opts = { + "model_type": params.get("model_type", "gguf"), + "custom_prompt": params.get("prompt", ""), + "skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true", + "keep_in_memory": str(params.get("keep_in_memory", "true")).lower() + == "true", + "general_threshold": params.get("general_threshold"), + "character_threshold": params.get("character_threshold"), + } + try: + for key in ("general_threshold", "character_threshold"): + if opts[key] is not None: + opts[key] = float(opts[key]) + except (ValueError, TypeError): + return self._bad_request( + "general_threshold and character_threshold must be numeric" + ) + + async def events(): + try: + from ..autotag import get_autotag_service + + async for event in self._autotag_batch_events( + get_autotag_service(), opts + ): + yield event + except Exception: + self.logger.exception("AutoTag error") + yield { + "type": "error", + "message": _INTERNAL_ERROR, + } + + if not _autotag_batch_lock.acquire(blocking=False): + return web.json_response( + {"success": False, "error": _BUSY_ERROR}, status=409 + ) + try: + return await self._stream_sse(request, events()) + finally: + _autotag_batch_lock.release() + + def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]: + """Return the real path when it is a media file inside an output directory.""" + if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS: + return None + for output_dir in self._get_all_output_dirs(): + if path_is_within(raw_path, output_dir): + real = os.path.realpath(raw_path) + return real if os.path.isfile(real) else None + return None + + @staticmethod + def _forbidden(raw_path: str): + return web.json_response( + { + "success": False, + "error": ( + f"{os.path.basename(str(raw_path).rstrip('/' + os.sep))} is not " + "an image inside the configured output directories" + ), + }, + status=403, + ) + + async def _select_autotag_target(self, data: dict): + """Pick the image to tag from an autotag/single body. + + Returns ``(image_path, prompt_id, None)`` on success or + ``(None, None, error_response)`` when the request must be rejected. + """ + image_id = data.get("image_id") + if image_id is not None: + if isinstance(image_id, bool) or not isinstance(image_id, (int, str)): + return None, None, self._bad_request("image_id must be an integer") + try: + image_id = int(image_id) + except (TypeError, ValueError): + return None, None, self._bad_request("image_id must be an integer") + record = await self._run_in_executor(self.db.get_image_by_id, image_id) + if not record or not record.get("image_path"): + return ( + None, + None, + web.json_response( + {"success": False, "error": f"Image {image_id} not found"}, + status=404, + ), + ) + stored_path = record["image_path"] + resolved = self._resolve_allowed_image_path(stored_path) + if resolved is None: + return None, None, self._forbidden(stored_path) + return resolved, record.get("prompt_id"), None + + raw_path = data.get("path") or data.get("image_path") + if not raw_path or not isinstance(raw_path, str): + return None, None, self._bad_request("image_id or path is required") + + resolved = self._resolve_allowed_image_path(raw_path) + if resolved is None: + return None, None, self._forbidden(raw_path) + + prompt_id = None + try: + prompt_id = await self._run_in_executor( + self.db.get_prompt_id_for_image, raw_path + ) + if prompt_id is None and resolved != raw_path: + prompt_id = await self._run_in_executor( + self.db.get_prompt_id_for_image, resolved + ) + except Exception as e: + self.logger.warning(f"Could not find linked prompt: {e}") + return resolved, prompt_id, None + + @staticmethod + def _bad_request(message: str): + return web.json_response({"success": False, "error": message}, status=400) + async def autotag_single(self, request): - """Generate tags for a single image.""" + """Generate tags for a single image selected by image_id or gallery path.""" try: data = await request.json() - image_path = data.get("image_path") + if not isinstance(data, dict): + return self._bad_request("Request body must be a JSON object") model_type = data.get("model_type", "gguf") custom_prompt = data.get("prompt") use_gpu = data.get("use_gpu", True) @@ -342,18 +522,13 @@ class AutotagRoutesMixin: if character_threshold is not None: character_threshold = float(character_threshold) except (ValueError, TypeError): - return web.json_response( - { - "success": False, - "error": "general_threshold and character_threshold must be numeric", - }, - status=400, + return self._bad_request( + "general_threshold and character_threshold must be numeric" ) - if not image_path: - return web.json_response( - {"success": False, "error": "image_path is required"}, status=400 - ) + image_path, prompt_id, error = await self._select_autotag_target(data) + if error is not None: + return error from ..autotag import get_autotag_service @@ -363,44 +538,34 @@ class AutotagRoutesMixin: not service.is_model_loaded() or service.get_loaded_model_type() != model_type ): - loop = asyncio.get_event_loop() - await loop.run_in_executor( - None, lambda: service.load_model(model_type, use_gpu) - ) + await self._run_in_executor(service.load_model, model_type, use_gpu) if custom_prompt: service.custom_prompt = custom_prompt - loop = asyncio.get_event_loop() - tags = await loop.run_in_executor( - None, - lambda: service.generate_tags( - image_path, - general_threshold=general_threshold, - character_threshold=character_threshold, - ), + tags = await self._run_in_executor( + service.generate_tags, + image_path, + general_threshold=general_threshold, + character_threshold=character_threshold, ) - prompt_id = None - try: - prompt_id = await self._run_in_executor( - self.db.get_prompt_id_for_image, image_path - ) - except Exception as e: - self.logger.warning(f"Could not find linked prompt: {e}") - return web.json_response( { "success": True, "tags": tags, "prompt_id": prompt_id, - "image_path": image_path, + "filename": os.path.basename(image_path), } ) + except json.JSONDecodeError: + return self._bad_request("Request body must be valid JSON") except Exception as e: self.logger.error(f"AutoTag single error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) async def apply_autotag(self, request): """Apply selected tags to a prompt.""" @@ -445,7 +610,9 @@ class AutotagRoutesMixin: except Exception as e: self.logger.error(f"Apply autotag error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) async def unload_autotag_model(self, request): """Manually unload the AutoTag model from memory.""" @@ -472,7 +639,37 @@ class AutotagRoutesMixin: except Exception as e: self.logger.error(f"Unload autotag model error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) + + @staticmethod + def _collect_scan_files(output_path: Path): + """Media files under ``output_path`` from a bounded, symlink-free walk. + + Descends at most ``_SCAN_MAX_DEPTH`` levels, skips ``thumbnails`` + trees and symlinked files, and stops after ``_SCAN_MAX_FILES`` files. + Returns ``(paths, truncated)``. + """ + found = [] + base_depth = len(output_path.parts) + for dirpath, dirnames, filenames in os.walk(output_path): + depth = len(Path(dirpath).parts) - base_depth + dirnames[:] = sorted( + name + for name in dirnames + if name != "thumbnails" and depth + 1 <= _SCAN_MAX_DEPTH + ) + for name in sorted(filenames): + if Path(name).suffix.lower() not in _MEDIA_EXTENSIONS: + continue + candidate = Path(dirpath) / name + if candidate.is_symlink(): + continue + if len(found) >= _SCAN_MAX_FILES: + return found, True + found.append(candidate) + return found, False async def scan_output_dir(self, request): """Scan ComfyUI output directory for images.""" @@ -485,63 +682,49 @@ class AutotagRoutesMixin: ) output_path = Path(output_dir) - image_extensions = [ - ".png", - ".jpg", - ".jpeg", - ".gif", - ".webp", - ".bmp", - ".tiff", - ] + files, truncated = await self._run_in_executor( + self._collect_scan_files, output_path + ) + has_thumbnails = (output_path / "thumbnails").is_dir() images = [] - seen_paths = set() - - for ext in image_extensions: - for pattern in [f"*{ext.lower()}", f"*{ext.upper()}"]: - for image_path in output_path.rglob(pattern): - if "thumbnails" not in image_path.parts: - normalized_path = str(image_path).lower() - if normalized_path not in seen_paths: - seen_paths.add(normalized_path) - - rel_path = image_path.relative_to(output_path) - - thumbnail_url = None - thumbnails_dir = output_path / "thumbnails" - if thumbnails_dir.exists(): - rel_path_no_ext = rel_path.with_suffix("") - thumbnail_rel_path = f"thumbnails/{rel_path_no_ext.as_posix()}_thumb{image_path.suffix}" - thumbnail_abs_path = ( - thumbnails_dir - / f"{rel_path_no_ext.as_posix()}_thumb{image_path.suffix}" - ) - if thumbnail_abs_path.exists(): - from urllib.parse import quote - - thumbnail_url = f'/prompt_manager/images/serve/{quote(thumbnail_rel_path, safe="/")}' - - from urllib.parse import quote as url_quote - - images.append( - { - "filename": image_path.name, - "path": str(image_path), - "relative_path": str(rel_path), - "url": f'/prompt_manager/images/serve/{url_quote(rel_path.as_posix(), safe="/")}', - "thumbnail_url": thumbnail_url, - } - ) + for image_path in files: + rel_path = image_path.relative_to(output_path) + thumbnail_url = None + if has_thumbnails: + thumb_rel = ( + f"thumbnails/{rel_path.with_suffix('').as_posix()}" + f"_thumb{image_path.suffix}" + ) + thumbnail_url = thumbnail_url_for( + thumb_rel, output_path / thumb_rel + ) + images.append( + { + "filename": image_path.name, + "path": self._public_path(image_path), + "relative_path": rel_path.as_posix(), + "url": "/prompt_manager/images/serve/" + + quote(rel_path.as_posix(), safe="/"), + "thumbnail_url": thumbnail_url, + } + ) images.sort(key=lambda x: x["filename"]) self.logger.info(f"Found {len(images)} images in output directory") return web.json_response( - {"success": True, "images": images, "count": len(images)} + { + "success": True, + "images": images, + "count": len(images), + "truncated": truncated, + } ) except Exception as e: self.logger.error(f"Scan output dir error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) diff --git a/py/api/images.py b/py/api/images.py index 320a583..d03e27d 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -1,15 +1,341 @@ """Image and gallery API routes for PromptManager.""" import asyncio +import hashlib import json import os import re +import threading import time as _time import urllib.parse from pathlib import Path from aiohttp import web -from PIL import Image +from PIL import Image, UnidentifiedImageError + +from .prompts import bad_request, parse_page_params, publish_image_paths +from .prompts import thumbnail_url_for +from .prompts import safe_error_message as _safe_error + +try: + from ...utils.parallel import map_parallel +except ImportError: + from utils.parallel import map_parallel + +# Only these file types are ever served by the image routes, regardless of +# what sits inside an allowed directory (defence in depth for the gallery). +IMAGE_EXTENSIONS = frozenset( + { + ".png", + ".jpg", + ".jpeg", + ".webp", + ".gif", + ".bmp", + ".tiff", + ".tif", + ".mp4", + ".webm", + ".mov", + } +) + + +def _public_path(path): + """Response-safe form of *path* (relative to the ComfyUI tree). + + Imported lazily: the package defines it after importing this module. + """ + from . import _public_path as public_path + + return public_path(path) + + +def _canonical(path): + """Resolve symlinks and normalise case so paths compare on every OS.""" + return Path(os.path.normcase(os.path.realpath(str(path)))) + + +def _is_within(path, root): + """True when *path* (after realpath) lives under *root* (after realpath).""" + return _canonical(path).is_relative_to(_canonical(root)) + + +def _has_traversal(relative): + """True for absolute paths, drive-qualified paths or any ``..`` segment.""" + if not relative or os.path.isabs(relative) or os.path.splitdrive(relative)[0]: + return True + if relative[0] in ("/", "\\"): + return True + return ".." in re.split(r"[\\/]+", relative) + + +def _forbidden(message="Access denied"): + return web.json_response({"success": False, "error": message}, status=403) + + +def _media_access_error(path, allowed_dirs): + """Return a 403 response when *path* must not be served, else None. + + Fails closed: an empty *allowed_dirs* denies everything. + """ + if Path(path).suffix.lower() not in SERVABLE_EXTENSIONS: + return _forbidden("Only media files can be served") + if not allowed_dirs: + return _forbidden("No allowed image directories configured") + if not any(_is_within(path, root) for root in allowed_dirs): + return _forbidden() + return None + + +# Media types the gallery scans and thumbnails (a superset is never served). +GALLERY_IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".webp", ".gif"}) +VIDEO_EXTENSIONS = frozenset({".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"}) +MEDIA_EXTENSIONS = GALLERY_IMAGE_EXTENSIONS | VIDEO_EXTENSIONS +# What the serve routes hand out: every image type above plus every video type +# the gallery lists, so a card never points at a 403. +SERVABLE_EXTENSIONS = IMAGE_EXTENSIONS | VIDEO_EXTENSIONS + +# Caps on request-triggered filesystem work. +MAX_SCAN_DEPTH = 12 +MAX_THUMBNAIL_FILES = 10_000 + + +def _worker_threads(): + """Configured thread count for thumbnail work. + + Imported lazily because py.config imports this package. + """ + from ..config import PromptManagerConfig + + return PromptManagerConfig.WORKER_THREADS + + +def _thumbnail_batch_size(workers): + """Files per executor job: enough to keep every worker thread busy.""" + return max(4, workers * 2) + + +def _temp_sibling(dst): + """Hidden temp path next to *dst* that keeps its extension (for format sniffing).""" + dst = Path(dst) + token = f"{os.getpid()}-{threading.get_ident()}" + return dst.with_name(f".{dst.stem}.part-{token}{dst.suffix}") + + +def _atomic_save(img, dst, fmt=None, **save_kwargs): + """Save a PIL image to *dst* without a partially written file ever being + visible at that path: write a temp sibling, then rename over it.""" + tmp = _temp_sibling(dst) + try: + img.save(tmp, fmt, **save_kwargs) + os.replace(tmp, dst) + finally: + if tmp.exists(): + tmp.unlink() + + +def _write_image_thumbnail(src, dst, thumbnail_size): + """Resize *src* into *dst* (blocking PIL work), atomically.""" + with Image.open(src) as img: + if img.mode in ("RGBA", "LA", "P"): + img = img.convert("RGB") + img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS) + save_kwargs = {"quality": 85, "optimize": True} + if dst.suffix.lower() == ".png": + save_kwargs = {"optimize": True} + _atomic_save(img, dst, **save_kwargs) + + +def _new_thumbnail_stats(): + return {"generated": 0, "skipped": 0, "errors": []} + + +def _record_thumbnail_result(stats, result): + action = result["action"] + if action == "generated": + stats["generated"] += 1 + elif action == "skipped": + stats["skipped"] += 1 + else: + stats["errors"].append(f"{result['file']}: {result.get('error', 'failed')}") + + +def _thumbnail_progress_payload(index, total, stats, result, start): + elapsed = _time.monotonic() - start + rate = index / elapsed if elapsed > 0 else 0 + eta = (total - index) / rate if rate > 0 else 0 + return { + "processed": index, + "total_images": total, + "generated": stats["generated"], + "skipped": stats["skipped"], + "error_count": len(stats["errors"]), + "percentage": round(index / total * 100, 1) if total else 100.0, + "rate": round(rate, 1), + "eta": round(eta), + "elapsed": round(elapsed, 1), + "current_file": f"{result['dir']}/{result['file']}", + "file_type": result["type"], + "action": result["action"], + } + + +def _thumbnail_complete_payload(total, stats, start, remaining=0, processed=None): + """Final payload of a thumbnail run. + + ``total`` counts every media file considered (pending, already fresh and + left over), ``processed`` the files this run actually worked on (defaults + to ``total``) and ``remaining`` the pending files left for the next run + because of MAX_THUMBNAIL_FILES. + """ + elapsed = _time.monotonic() - start + processed = total if processed is None else processed + errors = stats["errors"] + message = ( + f"Generated {stats['generated']} new thumbnails, " + f"skipped {stats['skipped']} existing" + ) + if remaining: + message += f", {remaining} more to generate on the next run" + if errors: + message += f" ({len(errors)} errors occurred)" + return { + "count": stats["generated"], + "skipped": stats["skipped"], + "total_images": total, + "remaining": remaining, + "errors": errors[:10], + "error_count": len(errors), + "elapsed_time": round(elapsed, 2), + "processing_rate": round(processed / elapsed if elapsed > 0 else 0, 2), + "message": message, + } + + +def _thumbnail_is_fresh(src, dst): + """True when ``dst`` exists and is newer than ``src`` (blocking stat).""" + try: + return dst.stat().st_mtime > src.stat().st_mtime + except OSError: + return False + + +# Caps on request-driven work (list sizes, scans, request bodies). +MAX_GALLERY_FILES = 50_000 +MAX_BULK_LIMIT = 5_000 +MAX_JSON_BODY_BYTES = 1024 * 1024 +_READ_CHUNK = 64 * 1024 + + +class _BodyTooLarge(Exception): + """Raised when a request body exceeds MAX_JSON_BODY_BYTES.""" + + +def _json_error(message, status): + return web.json_response({"success": False, "error": message}, status=status) + + +async def _read_json_body(request): + """Read a JSON object body in bounded chunks. + + Raises _BodyTooLarge past MAX_JSON_BODY_BYTES (declared or streamed) and + ValueError when the body is not a JSON object. + """ + limit = MAX_JSON_BODY_BYTES + if request.content_length is not None and request.content_length > limit: + raise _BodyTooLarge() + chunks, size = [], 0 + while True: + chunk = await request.content.read(_READ_CHUNK) + if not chunk: + break + size += len(chunk) + if size > limit: + raise _BodyTooLarge() + chunks.append(chunk) + data = json.loads(b"".join(chunks).decode("utf-8")) + if not isinstance(data, dict): + raise ValueError("JSON body must be an object") + return data + + +async def _json_body_or_error(request): + """Return ``(data, None)`` or ``(None, error_response)``.""" + try: + return await _read_json_body(request), None + except _BodyTooLarge: + return None, _json_error("Request body too large", 413) + except (ValueError, UnicodeDecodeError): + return None, _json_error("Invalid JSON body", 400) + + +def _iter_media_files(output_path, max_files): + """Yield media files under *output_path* with bounded effort. + + Symlinks are never followed, ``thumbnails/`` is skipped, descent stops + at MAX_SCAN_DEPTH and at most *max_files* paths are produced. + """ + output_path = Path(output_path) + root_depth = len(output_path.parts) + produced = 0 + for root, dirs, files in os.walk(output_path): + root_path = Path(root) + rel_parts = root_path.parts[root_depth:] + if "thumbnails" in rel_parts: + dirs[:] = [] + continue + if len(rel_parts) >= MAX_SCAN_DEPTH: + dirs[:] = [] + for name in sorted(files): + if Path(name).suffix.lower() not in MEDIA_EXTENSIONS: + continue + if (root_path / name).is_symlink(): + continue # never read through a link out of the output tree + if produced >= max_files: + return + produced += 1 + yield root_path / name + + +def _output_image_entry(media_path, output_path, root_index, public_root=None): + """Build one gallery entry for *media_path* (blocking stat). + + ``path`` is relative to its root (what delete/autotag resolve against) and + ``root_dir`` is the root's public form; neither reveals the server layout. + """ + if public_root is None: + public_root = _public_path(output_path) + stat = media_path.stat() + rel_path = media_path.relative_to(output_path) + extension = media_path.suffix.lower() + is_video = extension in VIDEO_EXTENSIONS + + thumb_rel = ( + f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb" + f"{'.jpg' if is_video else extension}" + ) + thumbnail_url = thumbnail_url_for(thumb_rel, output_path / thumb_rel) + + return { + "id": hashlib.sha1(str(media_path).encode("utf-8")).hexdigest()[:16], + "filename": media_path.name, + "path": rel_path.as_posix(), + "relative_path": str(rel_path), + "root_dir": public_root, + "root_index": root_index, + "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", + "thumbnail_url": thumbnail_url, + "size": stat.st_size, + "modified_time": stat.st_mtime, + "extension": extension, + "media_type": "video" if is_video else "image", + "is_video": is_video, + } + + +def _in_subfolder(rel_dir, subfolder): + return rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep) class ImageRoutesMixin: @@ -60,7 +386,9 @@ class ImageRoutesMixin: async def generate_thumbnails_route(request): return await self.generate_thumbnails(request) - @routes.get("/prompt_manager/images/generate-thumbnails/progress") + # POST only: the stream writes files, so it must not be reachable + # from a prefetch, link preview or . + @routes.post("/prompt_manager/images/generate-thumbnails/progress") async def generate_thumbnails_progress_route(request): return await self.generate_thumbnails_with_progress(request) @@ -72,6 +400,16 @@ class ImageRoutesMixin: async def get_gallery_subfolders_route(request): return await self.get_gallery_subfolders(request) + # POST: drops the server-side listing cache and rewalks the output roots. + @routes.post("/prompt_manager/gallery/rescan") + async def rescan_gallery_route(request): + return await self.rescan_gallery(request) + + def _present_images(self, images): + """Image dicts ready for a response: urls added, server paths hidden.""" + self._enrich_images(images) + return publish_image_paths(images, _public_path) + async def get_prompt_images(self, request): """Get all images for a specific prompt.""" try: @@ -82,7 +420,7 @@ class ImageRoutesMixin: cleaned_images = [self._clean_nan_recursive(image) for image in images] # Add url and thumbnail_url so the frontend can serve them - self._enrich_images(cleaned_images) + cleaned_images = self._present_images(cleaned_images) # Additional fallback: convert to JSON string and clean NaN values manually try: @@ -108,30 +446,75 @@ class ImageRoutesMixin: except Exception as e: self.logger.error(f"Get prompt images error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def get_recent_images(self, request): - """Get recently generated images.""" + """Get recently generated images (bounded window, newest first).""" try: - limit = int(request.query.get("limit", 50)) - images = await self._run_in_executor(self.db.get_recent_images, limit) + try: + limit, offset = parse_page_params(request.query) + except ValueError: + return bad_request("limit and offset must be integers") - return web.json_response({"success": True, "images": images}) - except Exception as e: - self.logger.error(f"Get recent images error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) - - async def get_all_images(self, request): - """Get all generated images with linked prompts.""" - try: - images = await self._run_in_executor(self.db.get_all_images) + images = self._present_images( + await self._run_in_executor(self.db.get_recent_images, limit, offset) + ) return web.json_response( - {"success": True, "images": images, "count": len(images)} + { + "success": True, + "images": images, + "pagination": { + "limit": limit, + "offset": offset, + "count": len(images), + }, + } + ) + except Exception as e: + self.logger.error(f"Get recent images error: {e}") + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) + + async def get_all_images(self, request): + """Get generated images with linked prompts (bulk-capped pages).""" + try: + try: + limit, offset = parse_page_params( + request.query, + default_limit=MAX_BULK_LIMIT, + max_limit=MAX_BULK_LIMIT, + ) + except ValueError: + return bad_request("limit and offset must be integers") + + images = self._present_images( + await self._run_in_executor( + self.db.get_all_images, limit=limit, offset=offset + ) + ) + + return web.json_response( + { + "success": True, + "images": images, + "count": len(images), + "pagination": { + "limit": limit, + "offset": offset, + "count": len(images), + "has_more": len(images) == limit, + }, + } ) except Exception as e: self.logger.error(f"Get all images error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def search_images(self, request): """Search images by prompt text.""" @@ -142,41 +525,51 @@ class ImageRoutesMixin: {"success": False, "error": "Search query required"}, status=400 ) - images = await self._run_in_executor(self.db.search_images_by_prompt, query) + try: + limit, offset = parse_page_params(request.query) + except ValueError: + return bad_request("limit and offset must be integers") + + images = self._present_images( + await self._run_in_executor( + self.db.search_images_by_prompt, query, limit, offset + ) + ) return web.json_response( - {"success": True, "images": images, "query": query} + { + "success": True, + "images": images, + "query": query, + "pagination": { + "limit": limit, + "offset": offset, + "count": len(images), + }, + } ) except Exception as e: self.logger.error(f"Search images error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) def _scan_gallery_files_sync(self, output_path): """Scan output directory for media files (blocking I/O, run in executor). - Returns list of (path, mtime) tuples sorted by mtime descending. + Returns list of (path, mtime) tuples sorted by mtime descending; the + walk is bounded by MAX_SCAN_DEPTH and MAX_GALLERY_FILES. """ - image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] - video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] - media_extensions = image_extensions + video_extensions - all_images = [] - - seen_paths = set() - for ext in media_extensions: - for pattern in [f"*{ext}", f"*{ext.upper()}"]: - for media_path in output_path.rglob(pattern): - if "thumbnails" not in media_path.parts: - normalized_path = str(media_path).lower() - if normalized_path not in seen_paths: - seen_paths.add(normalized_path) - try: - mtime = media_path.stat().st_mtime - all_images.append((media_path, mtime)) - except OSError: - continue - - all_images.sort(key=lambda x: x[1], reverse=True) - return all_images + found = [] + for media_path in _iter_media_files(output_path, MAX_GALLERY_FILES): + try: + found.append((media_path, media_path.stat().st_mtime)) + except OSError: + continue + if len(found) >= MAX_GALLERY_FILES: + self.logger.warning(f"Gallery scan capped at {MAX_GALLERY_FILES} files") + found.sort(key=lambda x: x[1], reverse=True) + return found async def _get_gallery_files(self, output_path): """Get gallery files with per-directory TTL cache.""" @@ -193,8 +586,6 @@ class ImageRoutesMixin: async def get_output_images(self, request): """Get all images from ComfyUI output folder(s).""" try: - from urllib.parse import quote - output_dirs = self._get_all_output_dirs() if not output_dirs: return web.json_response( @@ -204,78 +595,16 @@ class ImageRoutesMixin: "images": [], }, ) - - limit = int(request.query.get("limit", 100)) - offset = int(request.query.get("offset", 0)) + try: + limit, offset = parse_page_params(request.query, default_limit=100) + except ValueError: + return bad_request("limit and offset must be integers") subfolder = request.query.get("subfolder", "").strip() - # Collect (path, mtime, root, root_index) from all directories - all_images = [] - for root_idx, output_path in enumerate(output_dirs): - dir_images = await self._get_gallery_files(output_path) - for img_path, mtime in dir_images: - all_images.append((img_path, mtime, output_path, root_idx)) - - # Sort combined results by mtime (newest first) β€” no .stat() calls - all_images.sort(key=lambda x: x[1], reverse=True) - - # Apply subfolder filter if provided - if subfolder: - filtered = [] - for img_path, mtime, root, ridx in all_images: - rel_dir = str(img_path.relative_to(root).parent) - if rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep): - filtered.append((img_path, mtime, root, ridx)) - all_images = filtered - + all_images = await self._collect_output_images(output_dirs, subfolder) total = len(all_images) - paginated = all_images[offset : offset + limit] - - video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] - - def _format_page(): - images = [] - for media_path, mtime, output_path, root_index in paginated: - try: - stat = media_path.stat() - rel_path = media_path.relative_to(output_path) - extension = media_path.suffix.lower() - is_video = extension in video_extensions - media_type = "video" if is_video else "image" - - thumbnail_url = None - thumbnails_dir = output_path / "thumbnails" - if thumbnails_dir.exists(): - thumbnail_ext = ".jpg" if is_video else extension - rel_path_no_ext = rel_path.with_suffix("") - thumbnail_rel_path = f"thumbnails/{rel_path_no_ext.as_posix()}_thumb{thumbnail_ext}" - thumbnail_abs_path = output_path / thumbnail_rel_path - if thumbnail_abs_path.exists(): - thumbnail_url = f'/prompt_manager/images/serve/{quote(thumbnail_rel_path, safe="/")}' - - images.append( - { - "id": str(hash(str(media_path))), - "filename": media_path.name, - "path": str(media_path), - "relative_path": str(rel_path), - "root_dir": str(output_path), - "root_index": root_index, - "url": f"/prompt_manager/images/serve/{rel_path.as_posix()}", - "thumbnail_url": thumbnail_url, - "size": stat.st_size, - "modified_time": stat.st_mtime, - "extension": extension, - "media_type": media_type, - "is_video": is_video, - } - ) - except Exception as e: - self.logger.error(f"Error processing media {media_path}: {e}") - continue - return images - - images = await self._run_in_executor(_format_page) + page = all_images[offset : offset + limit] + images = await self._run_in_executor(self._format_output_page, page) return web.json_response( { @@ -289,7 +618,9 @@ class ImageRoutesMixin: ) except Exception as e: self.logger.error(f"Output images error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def serve_image(self, request): """Serve the actual image file using streamed FileResponse.""" @@ -302,47 +633,19 @@ class ImageRoutesMixin: {"success": False, "error": "Image not found"}, status=404 ) - image_path = Path(image["image_path"]).resolve() + image_path = Path(image["image_path"]) + allowed_dirs = list(self._get_all_output_dirs()) + self._lora_image_dirs() - # Validate path is within any allowed directory - allowed_dirs = list(self._get_all_output_dirs()) + denied = _media_access_error(image_path, allowed_dirs) + if denied is not None: + return denied - # Also allow LoRA directories when integration is enabled - try: - from ..config import IntegrationConfig - - if IntegrationConfig.LORA_MANAGER_ENABLED: - from ..lora_utils import ( - find_lora_directories, - get_lora_image_cache_dir, - ) - - lora_dirs = find_lora_directories( - IntegrationConfig.LORA_MANAGER_PATH - ) - allowed_dirs.extend(Path(d) for d in lora_dirs) - allowed_dirs.append(get_lora_image_cache_dir()) - except Exception: - # LoRA integration is optional β€” skip if unavailable - pass - - if allowed_dirs: - allowed = any( - image_path.is_relative_to(d.resolve()) for d in allowed_dirs - ) - if not allowed: - return web.json_response( - {"success": False, "error": "Access denied"}, status=403 - ) - - if not image_path.exists(): + if not image_path.is_file(): return web.json_response( {"success": False, "error": "Image file not found"}, status=404 ) - response = web.FileResponse(image_path) - response.headers["Cache-Control"] = "public, max-age=3600" - return response + return self._file_response(image_path) except ValueError: return web.json_response( @@ -350,52 +653,117 @@ class ImageRoutesMixin: ) except Exception as e: self.logger.error(f"Serve image error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) + + def _lora_image_dirs(self): + """Extra directories served when the LoRA integration is enabled.""" + try: + from ..config import IntegrationConfig + + if not IntegrationConfig.LORA_MANAGER_ENABLED: + return [] + from ..lora_utils import find_lora_directories, get_lora_image_cache_dir + + lora_root = IntegrationConfig.LORA_MANAGER_PATH or "" + if lora_root.strip(): + # A configured path comes from config.json, which can be hand + # edited: only a LoraManager install under custom_nodes may + # widen the serving allow-list. An empty path keeps the + # auto-detection, which itself only scans custom_nodes. + from ..lora_utils import resolve_lora_manager_path + + lora_root = resolve_lora_manager_path(lora_root) + if not lora_root: + return [] + + dirs = [Path(d) for d in find_lora_directories(lora_root)] + dirs.append(Path(get_lora_image_cache_dir())) + return dirs + except Exception: + # LoRA integration is optional; never widen access on failure. + return [] + + @staticmethod + def _file_response(path): + response = web.FileResponse(path) + response.headers["Cache-Control"] = "public, max-age=3600" + return response + + @staticmethod + def _select_root(output_dirs, root_param): + """Narrow *output_dirs* to the ``?root=N`` entry when it is valid.""" + if root_param is None: + return output_dirs + try: + idx = int(root_param) + except ValueError: + return output_dirs + if 0 <= idx < len(output_dirs): + return [output_dirs[idx]] + return output_dirs async def serve_output_image(self, request): - """Serve image file directly from ComfyUI output folder using streamed FileResponse.""" + """Serve a media file from the ComfyUI output folder(s).""" try: filepath = request.match_info["filepath"] + if _has_traversal(filepath): + return _forbidden() + if Path(filepath).suffix.lower() not in SERVABLE_EXTENSIONS: + return _forbidden("Only media files can be served") - # Search all configured output directories - output_dirs = self._get_all_output_dirs() + output_dirs = self._select_root( + list(self._get_all_output_dirs()), request.query.get("root") + ) if not output_dirs: - return web.json_response( - {"success": False, "error": "ComfyUI output directory not found"}, - status=404, - ) + return _forbidden("No allowed image directories configured") - # If a root index is provided, use only that root to avoid - # path collisions when multiple roots share the same relative path. - root_idx = request.query.get("root") - if root_idx is not None: - try: - idx = int(root_idx) - if 0 <= idx < len(output_dirs): - output_dirs = [output_dirs[idx]] - except (ValueError, IndexError): - pass - - # Try each root directory until the file is found for output_path in output_dirs: - image_path = (output_path / filepath).resolve() - - # Security check: must be within this output directory - if not image_path.is_relative_to(output_path.resolve()): + image_path = output_path / filepath + if not image_path.exists(): continue + # A symlink that points outside the root is a hard deny, not + # a fall-through to the next root. + if not _is_within(image_path, output_path): + return _forbidden() + if image_path.is_file(): + return self._file_response(image_path) - if image_path.exists(): - response = web.FileResponse(image_path) - response.headers["Cache-Control"] = "public, max-age=3600" - return response - + # ComfyUI's cache middleware would otherwise let browsers cache this + # miss for an hour, hiding a thumbnail generated a moment later. return web.json_response( - {"success": False, "error": "Image file not found"}, status=404 + {"success": False, "error": "Image file not found"}, + status=404, + headers={"Cache-Control": "no-store"}, ) except Exception as e: self.logger.error(f"Serve output image error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) + + async def rescan_gallery(self, request): + """Forget the cached listing and rewalk every output root now. + + The listing is cached for ``_gallery_cache_ttl`` seconds; the gallery's + Rescan button needs files added or removed on disk to show up at once. + """ + try: + self.invalidate_gallery_cache() + output_dirs = list(self._get_all_output_dirs()) + total = 0 + for output_path in output_dirs: + total += len(await self._get_gallery_files(output_path)) + return web.json_response( + {"success": True, "total": total, "roots": len(output_dirs)} + ) + except Exception: + self.logger.exception("Gallery rescan error") + return web.json_response( + {"success": False, "error": "Failed to rescan the gallery"}, status=500 + ) async def get_gallery_subfolders(self, request): """Get distinct subfolders from gallery output directories.""" @@ -425,20 +793,28 @@ class ImageRoutesMixin: ) except Exception as e: self.logger.error(f"Subfolders error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) + + # Fixed quality presets: clients choose a name, never a raw pixel size. + THUMBNAIL_SIZES = {"low": (150, 150), "medium": (300, 300), "high": (600, 600)} + + @classmethod + def _thumbnail_size(cls, quality): + """Map a quality preset name to a pixel size; anything else is medium.""" + if not isinstance(quality, str): + return cls.THUMBNAIL_SIZES["medium"] + return cls.THUMBNAIL_SIZES.get(quality, cls.THUMBNAIL_SIZES["medium"]) async def generate_thumbnails(self, request): - """Generate thumbnails for all images and videos in the ComfyUI output directory.""" + """Generate every thumbnail in one executor job (POST, blocking).""" try: - # Get request parameters - data = await request.json() - quality = data.get("quality", "medium") + data, error = await _json_body_or_error(request) + if error is not None: + return error + thumbnail_size = self._thumbnail_size(data.get("quality", "medium")) - # Map quality to size - size_map = {"low": (150, 150), "medium": (300, 300), "high": (600, 600)} - thumbnail_size = size_map.get(quality, (300, 300)) - - # Find ComfyUI output directory output_dir = self._find_comfyui_output_dir() if not output_dir: return web.json_response( @@ -447,479 +823,253 @@ class ImageRoutesMixin: ) output_path = Path(output_dir) - thumbnails_dir = output_path / "thumbnails" - - # Run the entire thumbnail generation in executor (heavy PIL I/O) result = await self._run_in_executor( self._generate_thumbnails_sync, output_path, - thumbnails_dir, + output_path / "thumbnails", thumbnail_size, ) - - # Invalidate gallery cache since thumbnails changed self.invalidate_gallery_cache() - return web.json_response(result) - except ImportError: - return web.json_response( - { - "success": False, - "error": "PIL (Pillow) library not available. Install with: pip install Pillow", - }, - status=500, - ) except Exception as e: self.logger.error(f"Generate thumbnails error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) def _generate_thumbnails_sync(self, output_path, thumbnails_dir, thumbnail_size): """Blocking thumbnail generation loop (run in executor).""" - import time - - thumbnails_dir.mkdir(exist_ok=True) - - self.logger.info("Scanning for media files to generate thumbnails...") - image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] - video_extensions = [".mp4", ".webm", ".avi", ".mov", ".mkv", ".m4v", ".wmv"] - media_extensions = image_extensions + video_extensions - media_files = [] - for root, dirs, files in os.walk(output_path): - if "thumbnails" in Path(root).parts: - continue - for file in files: - if any(file.lower().endswith(ext) for ext in media_extensions): - media_files.append(Path(root) / file) - - total_images = len(media_files) - self.logger.info(f"Found {total_images} media files to process for thumbnails") - - if total_images == 0: - return { - "success": True, - "count": 0, - "total_images": 0, - "message": "No media files found to process", - "errors": [], - } - - generated_count = 0 - skipped_count = 0 - errors = [] - start_time = time.time() - - for i, media_file in enumerate(media_files): - try: - rel_path = media_file.relative_to(output_path) - is_video = any( - media_file.name.lower().endswith(ext) for ext in video_extensions - ) - - rel_path_no_ext = rel_path.with_suffix("") - if is_video: - thumbnail_path = ( - thumbnails_dir / f"{rel_path_no_ext.as_posix()}_thumb.jpg" - ) - else: - thumbnail_path = ( - thumbnails_dir - / f"{rel_path_no_ext.as_posix()}_thumb{rel_path.suffix}" - ) - - if ( - thumbnail_path.exists() - and thumbnail_path.stat().st_mtime > media_file.stat().st_mtime - ): - skipped_count += 1 - continue - - thumbnail_path.parent.mkdir(parents=True, exist_ok=True) - - if is_video: - if self._generate_video_thumbnail( - media_file, thumbnail_path, thumbnail_size - ): - generated_count += 1 - else: - errors.append( - f"Failed to generate video thumbnail for {media_file.name}" - ) - else: - with Image.open(media_file) as img: - if img.mode in ("RGBA", "LA", "P"): - img = img.convert("RGB") - img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS) - save_kwargs = {"quality": 85, "optimize": True} - if thumbnail_path.suffix.lower() == ".png": - save_kwargs = {"optimize": True} - img.save(thumbnail_path, **save_kwargs) - generated_count += 1 - - if ( - generated_count % 100 == 0 - or (i + 1) % max(1, total_images // 10) == 0 - ): - progress = ((i + 1) / total_images) * 100 - elapsed = time.time() - start_time - rate = (i + 1) / elapsed if elapsed > 0 else 0 - eta = ((total_images - i - 1) / rate) if rate > 0 else 0 - self.logger.info( - f"Thumbnail progress: {i+1}/{total_images} ({progress:.1f}%) - " - f"Generated: {generated_count}, Skipped: {skipped_count}, " - f"Rate: {rate:.1f} img/s, ETA: {eta:.0f}s" - ) - - except Exception as e: - error_msg = ( - f"Failed to generate thumbnail for {media_file.name}: {str(e)}" - ) - errors.append(error_msg) - self.logger.warning(error_msg) - - elapsed_time = time.time() - start_time - self.logger.info( - f"Thumbnail generation completed: {generated_count} generated, " - f"{skipped_count} skipped, {len(errors)} errors in {elapsed_time:.1f}s" + start = _time.monotonic() + plan = self._plan_thumbnails(output_path, thumbnails_dir) + targets = plan["targets"] + stats = _new_thumbnail_stats() + stats["skipped"] = plan["skipped"] + for result in self._generate_batch(targets, thumbnail_size): + _record_thumbnail_result(stats, result) + payload = _thumbnail_complete_payload( + len(targets) + plan["skipped"] + plan["remaining"], + stats, + start, + remaining=plan["remaining"], + processed=len(targets), ) + payload["success"] = True + payload["thumbnails_path"] = _public_path(thumbnails_dir) + self.logger.info(payload["message"]) + return payload - return { - "success": True, - "count": generated_count, - "skipped": skipped_count, - "total_images": total_images, - "errors": errors, - "thumbnails_path": str(thumbnails_dir), - "elapsed_time": round(elapsed_time, 2), - "processing_rate": round( - (total_images / elapsed_time) if elapsed_time > 0 else 0, 2 - ), + def _plan_thumbnails(self, output_path, thumbnails_dir): + """Decide which media files need a thumbnail (blocking walk). + + Returns ``{"targets": [(src, dst, is_video)], "skipped": n, "remaining": n}``. + Files whose thumbnail is already newer than the source are counted as + skipped and never queued, so a large library that is mostly done only + pays for the files that changed. The walk is bounded like the gallery + listing (MAX_SCAN_DEPTH, MAX_GALLERY_FILES, no symlinks) and at most + MAX_THUMBNAIL_FILES pending files are queued per run; ``remaining`` + counts the pending files beyond that cap. + """ + output_path = Path(output_path) + targets = [] + skipped = 0 + remaining = 0 + for src in _iter_media_files(output_path, MAX_GALLERY_FILES): + suffix = src.suffix.lower() + is_video = suffix in VIDEO_EXTENSIONS + rel_no_ext = src.relative_to(output_path).with_suffix("") + thumb_suffix = ".jpg" if is_video else suffix + dst = thumbnails_dir / f"{rel_no_ext.as_posix()}_thumb{thumb_suffix}" + if not _is_within(dst, thumbnails_dir): + self.logger.warning(f"Skipping thumbnail outside safe dir: {src.name}") + continue + if _thumbnail_is_fresh(src, dst): + skipped += 1 + continue + if len(targets) >= MAX_THUMBNAIL_FILES: + remaining += 1 + continue + targets.append((src, dst, is_video)) + if remaining: + self.logger.warning( + f"Thumbnail run capped at {MAX_THUMBNAIL_FILES} new files; " + f"{remaining} more will be generated on the next run" + ) + return {"targets": targets, "skipped": skipped, "remaining": remaining} + + def _thumbnail_targets(self, output_path, thumbnails_dir): + """Pending (src, dst, is_video) triples; see _plan_thumbnails.""" + return self._plan_thumbnails(output_path, thumbnails_dir)["targets"] + + def _generate_batch(self, targets, thumbnail_size): + """Create thumbnails for ``targets`` on the configured worker threads. + + Blocking; returns one result dict per target, in order. + """ + + def one(target): + src, dst, is_video = target + return self._generate_one(src, dst, thumbnail_size, is_video) + + return map_parallel(one, targets, _worker_threads()) + + def _generate_one(self, src, dst, thumbnail_size, is_video): + """Create one thumbnail (blocking). Never raises; returns a result dict.""" + result = { + "file": src.name, + "dir": src.parent.name, + "type": "video" if is_video else "image", } + try: + if not src.is_file(): + return {**result, "action": "error", "error": "File no longer exists"} + if dst.exists() and dst.stat().st_mtime > src.stat().st_mtime: + return {**result, "action": "skipped"} + dst.parent.mkdir(parents=True, exist_ok=True) + if is_video: + if not self._generate_video_thumbnail(src, dst, thumbnail_size): + return {**result, "action": "error", "error": "Video decode failed"} + else: + _write_image_thumbnail(src, dst, thumbnail_size) + return {**result, "action": "generated"} + except UnidentifiedImageError: + return {**result, "action": "error", "error": "Not a valid image file"} + except Exception as e: + self.logger.warning(f"Thumbnail failed for {src.name}", exc_info=True) + return {**result, "action": "error", "error": _safe_error(e)} async def generate_thumbnails_with_progress(self, request): - """Generate thumbnails with Server-Sent Events progress updates.""" + """Generate thumbnails with Server-Sent Events progress updates (POST). + + ``quality`` is read from the query string so a POST-capable SSE + client can keep using the same URL shape. Every blocking step (scan, + PIL/ffmpeg) runs in the executor; the event loop only writes SSE + frames between files. + """ + thumbnail_size = self._thumbnail_size(request.query.get("quality", "medium")) + response = web.StreamResponse( + status=200, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await response.prepare(request) try: - import time + output_dir = self._find_comfyui_output_dir() + if not output_dir: + await self._emit_sse( + response, "error", {"error": "ComfyUI output directory not found"} + ) + return response - # Parse query parameters - quality = request.query.get("quality", "medium") - - # Map quality to size - size_map = {"low": (150, 150), "medium": (300, 300), "high": (600, 600)} - thumbnail_size = size_map.get(quality, (300, 300)) - - # Set up SSE response - response = web.StreamResponse( - status=200, - headers={ - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "Access-Control-Allow-Origin": "*", - }, + output_path = Path(output_dir) + await self._emit_sse( + response, + "status", + {"phase": "scanning", "message": "Scanning for images and videos..."}, ) - await response.prepare(request) - - async def send_progress(event_type, data): - """Send SSE event to client.""" - try: - message = f"event: {event_type}\ndata: {json.dumps(data)}\n\n" - await response.write(message.encode("utf-8")) - await asyncio.sleep(0.01) - except Exception as e: - self.logger.warning(f"Failed to send SSE message: {e}") - - try: - # Find ComfyUI output directory - output_dir = self._find_comfyui_output_dir() - if not output_dir: - await send_progress( - "error", {"error": "ComfyUI output directory not found"} - ) - return response - - output_path = Path(output_dir) - thumbnails_dir = output_path / "thumbnails" - thumbnails_dir.mkdir(exist_ok=True) - - # Send scanning event - await send_progress( - "status", - { - "phase": "scanning", - "message": f"Scanning {output_path} for images and videos to process...", - }, - ) - - self.logger.info( - f"Starting thumbnail generation scan in: {output_path}" - ) - - # Find all media files (images and videos) - image_extensions = [".png", ".jpg", ".jpeg", ".webp", ".gif"] - video_extensions = [ - ".mp4", - ".webm", - ".avi", - ".mov", - ".mkv", - ".m4v", - ".wmv", - ] - media_extensions = image_extensions + video_extensions - media_files = [] - scanned_dirs = 0 - - for root, dirs, files in os.walk(output_path): - if "thumbnails" in Path(root).parts: - continue - - scanned_dirs += 1 - if scanned_dirs % 5 == 0: - await send_progress( - "status", - { - "phase": "scanning", - "message": f"Scanning directories... ({scanned_dirs} checked, {len(media_files)} files found)", - }, - ) - - for file in files: - if any(file.lower().endswith(ext) for ext in media_extensions): - media_files.append(Path(root) / file) - - self.logger.info( - f"Scan complete: Found {len(media_files)} media files in {scanned_dirs} directories" - ) - - total_images = len(media_files) - - # Count images vs videos for more detail - image_count = sum( - 1 - for f in media_files - if any(f.name.lower().endswith(ext) for ext in image_extensions) - ) - video_count = total_images - image_count - - await send_progress( - "start", - { - "total_images": total_images, - "phase": "processing", - "message": f"Found {image_count} images and {video_count} videos to process", - "image_count": image_count, - "video_count": video_count, - }, - ) - - self.logger.info( - f"Starting thumbnail generation for {image_count} images and {video_count} videos" - ) - - if total_images == 0: - await send_progress( - "complete", - { - "count": 0, - "skipped": 0, - "total_images": 0, - "elapsed_time": 0, - "message": "No media files found to process", - }, - ) - return response - - generated_count = 0 - skipped_count = 0 - errors = [] - start_time = time.time() - - for i, media_file in enumerate(media_files): - try: - if not media_file.exists() or not media_file.is_file(): - continue - - is_video = any( - media_file.name.lower().endswith(ext) - for ext in video_extensions - ) - - rel_path = media_file.relative_to(output_path) - - rel_path_no_ext = rel_path.with_suffix("") - if is_video: - thumbnail_path = ( - thumbnails_dir - / f"{rel_path_no_ext.as_posix()}_thumb.jpg" - ) - else: - thumbnail_path = ( - thumbnails_dir - / f"{rel_path_no_ext.as_posix()}_thumb{rel_path.suffix}" - ) - - # Ensure thumbnail path is within our thumbnails directory - try: - thumbnail_path = thumbnail_path.resolve() - thumbnails_dir_resolved = thumbnails_dir.resolve() - if not str(thumbnail_path).startswith( - str(thumbnails_dir_resolved) - ): - self.logger.warning( - f"Skipping thumbnail outside safe directory: {thumbnail_path}" - ) - continue - except Exception as e: - self.logger.warning( - f"Path validation failed for {rel_path}: {e}" - ) - continue - - # Skip if thumbnail already exists and is newer than original - if ( - thumbnail_path.exists() - and thumbnail_path.stat().st_mtime - > media_file.stat().st_mtime - ): - skipped_count += 1 - else: - thumbnail_path.parent.mkdir(parents=True, exist_ok=True) - - if is_video: - if self._generate_video_thumbnail( - media_file, thumbnail_path, thumbnail_size - ): - generated_count += 1 - else: - errors.append( - f"Failed to generate video thumbnail for {media_file.name}" - ) - else: - with Image.open(media_file) as img: - if img.mode in ("RGBA", "LA", "P"): - img = img.convert("RGB") - img.thumbnail( - thumbnail_size, Image.Resampling.LANCZOS - ) - save_kwargs = {"quality": 85, "optimize": True} - if thumbnail_path.suffix.lower() == ".png": - save_kwargs = {"optimize": True} - img.save(thumbnail_path, **save_kwargs) - generated_count += 1 - - # Send progress update - if ( - (i + 1) % 5 == 0 - or (i + 1) % max(1, total_images // 100) == 0 - or i == total_images - 1 - ): - elapsed = time.time() - start_time - progress_percent = ((i + 1) / total_images) * 100 - rate = (i + 1) / elapsed if elapsed > 0 else 0 - eta = ((total_images - i - 1) / rate) if rate > 0 else 0 - - file_info = { - "name": media_file.name, - "dir": media_file.parent.name, - "type": "video" if is_video else "image", - "action": ( - "skipped" - if thumbnail_path.exists() - else "generating" - ), - } - - await send_progress( - "progress", - { - "processed": i + 1, - "total_images": total_images, - "generated": generated_count, - "skipped": skipped_count, - "percentage": round(progress_percent, 1), - "rate": round(rate, 1), - "eta": round(eta, 0), - "elapsed": round(elapsed, 1), - "current_file": f"{file_info['dir']}/{file_info['name']}", - "file_type": file_info["type"], - "action": file_info["action"], - }, - ) - - if (i + 1) % 50 == 0: - self.logger.info( - f"Thumbnail progress: {i+1}/{total_images} ({progress_percent:.1f}%) - Generated: {generated_count}, Skipped: {skipped_count}" - ) - - except Exception as e: - self.logger.warning( - f"Failed to generate thumbnail for {media_file.name}", - exc_info=True, - ) - error_msg = ( - f"Failed to generate thumbnail for {media_file.name}" - ) - errors.append(error_msg) - - if len(errors) <= 5: - await send_progress( - "status", - { - "phase": "processing", - "message": f"Error processing {media_file.name}", - }, - ) - - elapsed_time = time.time() - start_time - - completion_message = f"Successfully generated {generated_count} new thumbnails, skipped {skipped_count} existing" - if errors: - completion_message += f" ({len(errors)} errors occurred)" - - await send_progress( - "complete", - { - "count": generated_count, - "skipped": skipped_count, - "total_images": total_images, - "errors": errors[:10], - "error_count": len(errors), - "elapsed_time": round(elapsed_time, 2), - "processing_rate": round( - (total_images / elapsed_time) if elapsed_time > 0 else 0, 2 - ), - "message": completion_message, - }, - ) - - self.logger.info( - f"Thumbnail generation completed: {generated_count} generated, {skipped_count} skipped, {len(errors)} errors in {elapsed_time:.2f}s" - ) - - except Exception as e: - self.logger.exception("Thumbnail generation failed") - await send_progress( - "error", - { - "error": "An internal error occurred", - "message": "Thumbnail generation failed. Check server logs for details.", - }, - ) - - return response - - except ImportError: - return web.json_response( + plan = await self._run_in_executor( + self._plan_thumbnails, output_path, output_path / "thumbnails" + ) + await self._stream_thumbnails( + response, + plan["targets"], + thumbnail_size, + skipped=plan["skipped"], + remaining=plan["remaining"], + ) + self.invalidate_gallery_cache() + except Exception: + self.logger.exception("Thumbnail generation failed") + await self._emit_sse( + response, + "error", { - "success": False, - "error": "PIL (Pillow) library not available. Install with: pip install Pillow", + "error": "An internal error occurred", + "message": "Thumbnail generation failed. Check server logs.", }, - status=500, ) + return response + + async def _stream_thumbnails( + self, response, targets, thumbnail_size, skipped=0, remaining=0 + ): + """Process pending *targets* in worker-thread batches, emitting SSE per file. + + ``skipped`` files already had fresh thumbnails and ``remaining`` were + left for the next run; both only feed the start/complete payloads. + """ + total = len(targets) + video_count = sum(1 for _, _, is_video in targets if is_video) + message = ( + f"Found {total - video_count} images and {video_count} videos to process" + ) + if skipped: + message += f" ({skipped} already have thumbnails)" + connected = await self._emit_sse( + response, + "start", + { + "total_images": total, + "skipped": skipped, + "remaining": remaining, + "phase": "processing", + "image_count": total - video_count, + "video_count": video_count, + "message": message, + }, + ) + if not connected: + return + stats = _new_thumbnail_stats() + stats["skipped"] = skipped + start = _time.monotonic() + batch_size = _thumbnail_batch_size(_worker_threads()) + index = 0 + for batch_start in range(0, total, batch_size): + batch = targets[batch_start : batch_start + batch_size] + results = await self._run_in_executor( + self._generate_batch, batch, thumbnail_size + ) + for result in results: + index += 1 + _record_thumbnail_result(stats, result) + if result["action"] == "error": + await self._emit_sse( + response, + "file_error", + {"file": result["file"], "error": result["error"]}, + ) + connected = await self._emit_sse( + response, + "progress", + _thumbnail_progress_payload(index, total, stats, result, start), + ) + if not connected: + self.logger.info("Thumbnail client disconnected; stopping early") + return + await asyncio.sleep(0) # let other requests run between batches + payload = _thumbnail_complete_payload( + total + skipped + remaining, + stats, + start, + remaining=remaining, + processed=total, + ) + await self._emit_sse(response, "complete", payload) + self.logger.info(payload["message"]) + + async def _emit_sse(self, response, event, data) -> bool: + """Write one SSE frame. Returns False once the client has gone away.""" + try: + frame = f"event: {event}\ndata: {json.dumps(data)}\n\n" + await response.write(frame.encode("utf-8")) + return True except Exception as e: - self.logger.error(f"Generate thumbnails with progress error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + self.logger.warning(f"Failed to send SSE message: {e}") + return False def _generate_video_thumbnail(self, video_path, thumbnail_path, thumbnail_size): """Generate thumbnail from video file. Returns True if successful.""" @@ -946,7 +1096,7 @@ class ImageRoutesMixin: img = Image.fromarray(frame_rgb) img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS) - img.save(thumbnail_path, "JPEG", quality=85, optimize=True) + _atomic_save(img, thumbnail_path, "JPEG", quality=85, optimize=True) self.logger.debug( f"Generated video thumbnail using OpenCV: {thumbnail_path}" ) @@ -959,6 +1109,7 @@ class ImageRoutesMixin: try: import subprocess + tmp_path = _temp_sibling(thumbnail_path) cmd = [ "ffmpeg", "-i", @@ -970,11 +1121,12 @@ class ImageRoutesMixin: "-s", f"{thumbnail_size[0]}x{thumbnail_size[1]}", "-y", - str(thumbnail_path), + str(tmp_path), ] result = subprocess.run(cmd, capture_output=True, text=True, timeout=30) - if result.returncode == 0: + if result.returncode == 0 and tmp_path.exists(): + os.replace(tmp_path, thumbnail_path) self.logger.debug( f"Generated video thumbnail using ffmpeg: {thumbnail_path}" ) @@ -986,6 +1138,11 @@ class ImageRoutesMixin: except (ImportError, subprocess.TimeoutExpired, FileNotFoundError): pass + finally: + try: + tmp_path.unlink() + except (NameError, OSError): + pass # Last resort: create a placeholder thumbnail try: @@ -1021,7 +1178,7 @@ class ImageRoutesMixin: except (OSError, AttributeError): pass - img.save(thumbnail_path, "JPEG", quality=85) + _atomic_save(img, thumbnail_path, "JPEG", quality=85) self.logger.debug( f"Generated placeholder video thumbnail: {thumbnail_path}" ) @@ -1042,7 +1199,6 @@ class ImageRoutesMixin: async def clear_thumbnails(self, request): """Safely clear only our generated thumbnails, never touch original images.""" try: - import shutil # Find ComfyUI output directory output_dir = self._find_comfyui_output_dir() @@ -1064,24 +1220,23 @@ class ImageRoutesMixin: } ) - # Verify this is actually our thumbnails directory + # Verify this is actually our thumbnails directory: a real child + # of the output root (symlinks resolved), not merely a path that + # shares its prefix such as "2/thumbnails". try: - thumbnails_dir_resolved = thumbnails_dir.resolve() - output_path_resolved = output_path.resolve() - if ( - not str(thumbnails_dir_resolved).startswith( - str(output_path_resolved) - ) + not _is_within(thumbnails_dir, output_path) or thumbnails_dir.name != "thumbnails" ): self.logger.error( - f"Safety check failed: thumbnails directory path invalid: {thumbnails_dir}" + "Safety check failed: thumbnails directory path invalid: " + f"{thumbnails_dir}" ) return web.json_response( { "success": False, - "error": "Safety check failed: invalid thumbnails directory path", + "error": "Safety check failed: invalid thumbnails " + "directory path", }, status=400, ) @@ -1134,7 +1289,8 @@ class ImageRoutesMixin: return f"{bytes_size:.1f} TB" self.logger.info( - f"Thumbnail cleanup: cleared {cleared_count} files ({format_size(cleared_size)})" + f"Thumbnail cleanup: cleared {cleared_count} files " + f"({format_size(cleared_size)})" ) return web.json_response( @@ -1143,32 +1299,69 @@ class ImageRoutesMixin: "cleared_files": cleared_count, "cleared_size": cleared_size, "cleared_size_formatted": format_size(cleared_size), - "message": f"Cleared {cleared_count} thumbnail files ({format_size(cleared_size)})", + "message": f"Cleared {cleared_count} thumbnail files " + "({format_size(cleared_size)})", } ) except Exception as e: self.logger.error(f"Clear thumbnails error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) + + @staticmethod + def _resolve_client_media_path(raw_path, allowed_dirs): + """Absolute path for a client-supplied media path, or None when it + does not name a file inside an allowed root. + + Absolute paths are used as given (the caller still runs the access + check). Relative paths are the forms this API hands out: relative to + a gallery root (``sub/a.png``) or to a ComfyUI anchor + (``output/sub/a.png``), and must resolve inside an allowed root. + """ + if os.path.isabs(raw_path): + return raw_path + from ..config import GalleryConfig + + for base in [*allowed_dirs, *GalleryConfig.path_anchors()]: + candidate = Path(base) / raw_path + if candidate.is_file() and any( + _is_within(candidate, root) for root in allowed_dirs + ): + return str(candidate) + return None async def link_image_to_prompt(self, request): - """Link a generated image to a prompt.""" + """Link a generated image (inside an allowed directory) to a prompt.""" try: - data = await request.json() + data, error = await _json_body_or_error(request) + if error is not None: + return error prompt_id = data.get("prompt_id") image_path = data.get("image_path") metadata = data.get("metadata", {}) - if not prompt_id or not image_path: - return web.json_response( - { - "success": False, - "error": "prompt_id and image_path are required", - }, - status=400, - ) + if not prompt_id or not isinstance(image_path, str) or not image_path: + return bad_request("prompt_id and image_path are required") + if not os.path.isabs(image_path) and _has_traversal(image_path): + return _forbidden() - if not os.path.exists(image_path): + allowed_dirs = list(self._get_all_output_dirs()) + resolved = self._resolve_client_media_path(image_path, allowed_dirs) + if resolved is None: + if Path(image_path).suffix.lower() not in SERVABLE_EXTENSIONS: + return _forbidden("Only media files can be served") + return web.json_response( + {"success": False, "error": "Image file not found"}, status=404 + ) + image_path = resolved + + denied = _media_access_error(image_path, allowed_dirs) + if denied is not None: + return denied + + if not os.path.isfile(image_path): return web.json_response( {"success": False, "error": "Image file not found"}, status=404 ) @@ -1187,7 +1380,9 @@ class ImageRoutesMixin: except Exception as e: self.logger.error(f"Link image error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def get_image_prompt(self, request): """Get prompt information for a specific image path.""" @@ -1213,17 +1408,15 @@ class ImageRoutesMixin: self.db.get_image_prompt_info, image_path ) if prompt_data: - prompt_data["image_path"] = image_path - if prompt_data: + prompt_data["image_path"] = _public_path(image_path) return web.json_response({"success": True, "prompt": prompt_data}) - else: - return web.json_response( - { - "success": False, - "error": "No prompt found for this image", - "image_path": image_path, - } - ) + return web.json_response( + { + "success": False, + "error": "No prompt found for this image", + "image_path": _public_path(image_path), + } + ) except Exception as db_error: self.logger.error(f"Database error in get_image_prompt: {db_error}") @@ -1233,7 +1426,9 @@ class ImageRoutesMixin: except Exception as e: self.logger.error(f"Get image prompt error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) async def delete_image(self, request): """Delete an image record.""" @@ -1256,4 +1451,37 @@ class ImageRoutesMixin: ) except Exception as e: self.logger.error(f"Delete image error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": _safe_error(e)}, status=500 + ) + + async def _collect_output_images(self, output_dirs, subfolder): + """Merge cached per-root listings, newest first, optionally filtered.""" + merged = [] + for root_idx, output_path in enumerate(output_dirs): + for img_path, mtime in await self._get_gallery_files(output_path): + merged.append((img_path, mtime, output_path, root_idx)) + merged.sort(key=lambda x: x[1], reverse=True) + if not subfolder: + return merged + return [ + entry + for entry in merged + if _in_subfolder(str(entry[0].relative_to(entry[2]).parent), subfolder) + ] + + def _format_output_page(self, page): + entries = [] + public_roots = {} + for media_path, _mtime, output_path, root_index in page: + if output_path not in public_roots: + public_roots[output_path] = _public_path(output_path) + try: + entries.append( + _output_image_entry( + media_path, output_path, root_index, public_roots[output_path] + ) + ) + except Exception as e: + self.logger.error(f"Error processing media {media_path.name}: {e}") + return entries diff --git a/py/api/logging_routes.py b/py/api/logging_routes.py index 9121a86..5329a1f 100644 --- a/py/api/logging_routes.py +++ b/py/api/logging_routes.py @@ -4,6 +4,86 @@ import os from aiohttp import web +# Keys PromptManagerLogger.update_config() understands. Anything else is +# rejected so a client cannot inject arbitrary entries into the logging config. +LOG_CONFIG_KEYS = frozenset( + { + "level", + "max_file_size", + "backup_count", + "console_logging", + "file_logging", + "buffer_size", + } +) +_LOG_INT_KEYS = frozenset({"max_file_size", "backup_count", "buffer_size"}) +_LOG_BOOL_KEYS = frozenset({"console_logging", "file_logging"}) +_LOG_LEVEL_NAMES = frozenset({"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"}) +# Inclusive bounds for the integer keys that size the log files on disk +_MEGABYTE = 1024 * 1024 +_LOG_INT_BOUNDS = { + "max_file_size": (1 * _MEGABYTE, 1024 * _MEGABYTE), + "backup_count": (0, 50), +} +_LOG_INT_BOUND_LABELS = { + "max_file_size": "between 1 MB and 1024 MB (in bytes)", + "backup_count": "between 0 and 50", +} + +DEFAULT_LOG_LIMIT = 100 +MAX_LOG_LIMIT = 500 + + +def _parse_limit(raw): + """Parse a ``limit`` query value, clamped to [1, MAX_LOG_LIMIT]. + + Raises: + ValueError: when ``raw`` is not an integer. + """ + if raw is None or raw == "": + return DEFAULT_LOG_LIMIT + limit = int(raw) + return max(1, min(limit, MAX_LOG_LIMIT)) + + +def _validate_log_config(data): + """Validate a log-config payload against LOG_CONFIG_KEYS. + + Returns: + (clean, error): ``clean`` is the normalised dict to apply, ``error`` + a message describing the first problem (``clean`` is then ``None``). + """ + if not isinstance(data, dict): + return None, "Request body must be a JSON object" + + unknown = sorted(key for key in data if key not in LOG_CONFIG_KEYS) + if unknown: + return None, ( + f"Unknown logging config keys: {', '.join(unknown)}. " + f"Allowed: {', '.join(sorted(LOG_CONFIG_KEYS))}" + ) + + clean = {} + for key, value in data.items(): + if key == "level": + if not isinstance(value, str) or value.upper() not in _LOG_LEVEL_NAMES: + return None, ( + f"Invalid log level. Must be one of: {sorted(_LOG_LEVEL_NAMES)}" + ) + clean[key] = value.upper() + elif key in _LOG_INT_KEYS: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + return None, f"{key} must be a non-negative integer" + bounds = _LOG_INT_BOUNDS.get(key) + if bounds and not bounds[0] <= value <= bounds[1]: + return None, f"{key} must be {_LOG_INT_BOUND_LABELS[key]}" + clean[key] = value + elif key in _LOG_BOOL_KEYS: + if not isinstance(value, bool): + return None, f"{key} must be a boolean" + clean[key] = value + return clean, None + class LoggingRoutesMixin: """Mixin providing logging-related API endpoints.""" @@ -56,14 +136,15 @@ class LoggingRoutesMixin: try: logger_manager = self._get_logger_manager() - limit = int(request.query.get("limit", 100)) + try: + limit = _parse_limit(request.query.get("limit")) + except ValueError: + return web.json_response( + {"success": False, "error": "limit must be an integer", "logs": []}, + status=400, + ) level = request.query.get("level", None) - if limit > 1000: - limit = 1000 - elif limit < 1: - limit = 1 - logs = logger_manager.get_recent_logs(limit=limit, level=level) return web.json_response( @@ -79,14 +160,18 @@ class LoggingRoutesMixin: except Exception as e: self.logger.error(f"Get logs error: {e}") return web.json_response( - {"success": False, "error": str(e), "logs": []}, status=500 + {"success": False, "error": self._public_error(e), "logs": []}, + status=500, ) async def get_log_files(self, request): """Get information about available log files.""" try: logger_manager = self._get_logger_manager() - log_files = logger_manager.get_log_files() + log_files = [ + {**entry, "path": self._public_path(entry.get("path"))} + for entry in logger_manager.get_log_files() + ] return web.json_response( {"success": True, "files": log_files, "count": len(log_files)} @@ -95,7 +180,8 @@ class LoggingRoutesMixin: except Exception as e: self.logger.error(f"Get log files error: {e}") return web.json_response( - {"success": False, "error": str(e), "files": []}, status=500 + {"success": False, "error": self._public_error(e), "files": []}, + status=500, ) async def download_log_file(self, request): @@ -130,7 +216,9 @@ class LoggingRoutesMixin: except Exception as e: self.logger.error(f"Download log file error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) async def truncate_logs(self, request): """Truncate all log files.""" @@ -148,7 +236,9 @@ class LoggingRoutesMixin: except Exception as e: self.logger.error(f"Truncate logs error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) async def get_log_config(self, request): """Get current logging configuration.""" @@ -160,27 +250,23 @@ class LoggingRoutesMixin: except Exception as e: self.logger.error(f"Get log config error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) async def update_log_config(self, request): - """Update logging configuration.""" + """Update logging configuration (whitelisted keys only).""" try: - data = await request.json() + data, error_response = await self._read_json_body(request) + if error_response is not None: + return error_response + + clean, error = _validate_log_config(data) + if error: + return web.json_response({"success": False, "error": error}, status=400) + logger_manager = self._get_logger_manager() - - if "level" in data: - valid_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] - if data["level"].upper() not in valid_levels: - return web.json_response( - { - "success": False, - "error": f"Invalid log level. Must be one of: {valid_levels}", - }, - status=400, - ) - data["level"] = data["level"].upper() - - logger_manager.update_config(data) + logger_manager.update_config(clean) return web.json_response( { @@ -192,16 +278,22 @@ class LoggingRoutesMixin: except Exception as e: self.logger.error(f"Update log config error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": "Failed to update logging configuration"}, + status=500, + ) async def get_log_stats(self, request): """Get logging statistics.""" try: logger_manager = self._get_logger_manager() - stats = logger_manager.get_log_stats() + stats = dict(logger_manager.get_log_stats()) + stats["log_directory"] = self._public_path(stats.get("log_directory")) return web.json_response({"success": True, "stats": stats}) except Exception as e: self.logger.error(f"Get log stats error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) diff --git a/py/api/lora_integration.py b/py/api/lora_integration.py index 6275bfd..537b35c 100644 --- a/py/api/lora_integration.py +++ b/py/api/lora_integration.py @@ -1,11 +1,15 @@ """LoraManager integration API routes for PromptManager.""" import json -import os from pathlib import Path from aiohttp import web +_LORA_PATH_ERROR = ( + "LoraManager path must be a directory under ComfyUI's custom_nodes " + "whose name contains 'lora'" +) + class LoraIntegrationMixin: """Mixin providing LoraManager detection, scanning, and trigger word endpoints.""" @@ -35,6 +39,15 @@ class LoraIntegrationMixin: async def lora_refresh_cache_route(request): return await self.lora_refresh_cache(request) + def _error_response(self, exc: Exception, status: int = 500): + return web.json_response( + {"success": False, "error": self._public_error(exc)}, status=status + ) + + @staticmethod + def _bad_lora_request(message: str): + return web.json_response({"success": False, "error": message}, status=400) + # ── Detection ──────────────────────────────────────────────────── async def lora_detect(self, request): @@ -47,12 +60,12 @@ class LoraIntegrationMixin: { "success": True, "detected": path is not None, - "path": path or "", + "path": self._public_path(path) if path else "", } ) except Exception as e: self.logger.error(f"LoraManager detection failed: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return self._error_response(e) # ── Status ─────────────────────────────────────────────────────── @@ -74,20 +87,49 @@ class LoraIntegrationMixin: { "success": True, "enabled": config["enabled"], - "path": config["path"], + "path": self._public_path(config["path"]), "trigger_words_enabled": config["trigger_words_enabled"], - "civitai_api_key": config.get("civitai_api_key", ""), + "has_civitai_api_key": bool(config.get("civitai_api_key")), "detected": detected_path is not None, - "detected_path": detected_path or "", + "detected_path": ( + self._public_path(detected_path) if detected_path else "" + ), "trigger_cache_loaded": cache.is_loaded, } ) except Exception as e: self.logger.error(f"LoraManager status check failed: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return self._error_response(e) # ── Enable / Disable ───────────────────────────────────────────── + async def _validated_lora_manager_path(self, path): + """Canonical LoraManager path for ``path`` or ``(None, error_response)``. + + The user's path (or, when empty, the auto-detected install) must be a + LoraManager-looking directory that is a direct child of ComfyUI's + ``custom_nodes`` with ``lora`` in its name; the trees under it are + later served as images, so nothing outside custom_nodes may be used. + """ + from ..lora_utils import detect_lora_manager, resolve_lora_manager_path + + candidate = "" + if path: + candidate = await self._run_in_executor(resolve_lora_manager_path, path) + if candidate is None: + return None, self._bad_lora_request(_LORA_PATH_ERROR) + + resolved = await self._run_in_executor(detect_lora_manager, candidate) + if not resolved: + return None, self._bad_lora_request( + "LoraManager not found at the specified path" + ) + + allowed = await self._run_in_executor(resolve_lora_manager_path, resolved) + if allowed is None: + return None, self._bad_lora_request(_LORA_PATH_ERROR) + return allowed, None + async def lora_enable(self, request): """Enable or disable LoraManager integration and save to config.json.""" try: @@ -95,23 +137,29 @@ class LoraIntegrationMixin: enabled = data.get("enabled", False) path = data.get("path", "") trigger_words = data.get("trigger_words_enabled", False) - civitai_key = data.get("civitai_api_key", "") from ..config import IntegrationConfig, PromptManagerConfig - from ..lora_utils import detect_lora_manager, get_trigger_cache + from ..lora_utils import get_trigger_cache + + # The key is never echoed back to the UI, so a blank or absent + # field means "unchanged"; clearing is an explicit request. + civitai_key = IntegrationConfig.CIVITAI_API_KEY + submitted_key = data.get("civitai_api_key") + if submitted_key is not None and not isinstance(submitted_key, str): + return self._bad_lora_request("civitai_api_key must be a string") + if data.get("clear_civitai_api_key"): + civitai_key = "" + elif submitted_key and submitted_key.strip(): + civitai_key = submitted_key.strip() + + if not isinstance(path, str): + return self._bad_lora_request("path must be a string") # If enabling, validate the path if enabled: - resolved = await self._run_in_executor(detect_lora_manager, path) - if not resolved: - return web.json_response( - { - "success": False, - "error": "LoraManager not found at the specified path", - }, - status=400, - ) - path = resolved + path, error = await self._validated_lora_manager_path(path) + if error is not None: + return error # Update in-memory config IntegrationConfig.LORA_MANAGER_ENABLED = enabled @@ -119,12 +167,8 @@ class LoraIntegrationMixin: IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = trigger_words IntegrationConfig.CIVITAI_API_KEY = civitai_key - # Persist to config.json - config_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) - config_file = os.path.join(config_dir, "config.json") - PromptManagerConfig.save_to_file(config_file) + # Persist to the configured config.json + PromptManagerConfig.save_to_file() # Load trigger word cache if enabling cache = get_trigger_cache() @@ -138,49 +182,36 @@ class LoraIntegrationMixin: { "success": True, "enabled": enabled, - "path": path, + "path": self._public_path(path), "trigger_words_enabled": trigger_words, } ) except Exception as e: self.logger.error(f"LoraManager enable/disable failed: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return self._error_response(e) # ── Scan LoRA example images ───────────────────────────────────── async def lora_scan(self, request): """Scan LoraManager metadata and import LoRA info + preview images. - Streams progress as SSE, matching the existing scan pattern. + Streams progress as SSE, matching the existing scan pattern. Errors + raised once the stream is open are reported as an ``error`` event. """ try: from ..config import IntegrationConfig - from ..lora_utils import ( - download_civitai_images, - find_lora_directories, - get_example_prompt_from_metadata, - get_lora_image_cache_dir, - get_preview_images_from_metadata, - get_trigger_words_from_metadata, - get_model_name_from_metadata, - read_lora_metadata, - ) + from ..lora_utils import resolve_lora_manager_path if not IntegrationConfig.LORA_MANAGER_ENABLED: - return web.json_response( - { - "success": False, - "error": "LoraManager integration is not enabled", - }, - status=400, - ) + return self._bad_lora_request("LoraManager integration is not enabled") lm_path = IntegrationConfig.LORA_MANAGER_PATH if not lm_path: - return web.json_response( - {"success": False, "error": "LoraManager path not configured"}, - status=400, - ) + return self._bad_lora_request("LoraManager path not configured") + + allowed = await self._run_in_executor(resolve_lora_manager_path, lm_path) + if allowed is None: + return self._bad_lora_request(_LORA_PATH_ERROR) response = web.StreamResponse( status=200, @@ -193,167 +224,185 @@ class LoraIntegrationMixin: line = f"data: {json.dumps(data)}\n\n" await response.write(line.encode("utf-8")) - await send_progress( - { - "type": "progress", - "status": "Clearing previous lora-manager imports...", - "progress": 0, - } - ) - - # Clear previous imports so reimport is always clean - await self._run_in_executor( - self.db.delete_prompts_by_category, "lora-manager" - ) - - await send_progress( - { - "type": "progress", - "status": "Finding LoRA directories...", - "progress": 2, - } - ) - - lora_dirs = await self._run_in_executor(find_lora_directories, lm_path) - - # Collect all metadata files - meta_files = [] - for d in lora_dirs: - dir_path = Path(d) - meta_files.extend(dir_path.rglob("*.metadata.json")) - - total = len(meta_files) - imported = 0 - skipped = 0 - - await send_progress( - { - "type": "progress", - "status": f"Found {total} LoRA metadata files", - "progress": 5, - "total": total, - } - ) - - cache_dir = get_lora_image_cache_dir() - - for i, meta_file in enumerate(meta_files): - metadata = await self._run_in_executor(read_lora_metadata, meta_file) - if not metadata: - skipped += 1 - continue - - model_name = get_model_name_from_metadata(metadata) - trigger_words = get_trigger_words_from_metadata(metadata) - - # Collect all images: local previews + downloaded civitai examples - preview_paths = await self._run_in_executor( - get_preview_images_from_metadata, metadata, meta_file - ) - civitai_paths = await self._run_in_executor( - download_civitai_images, - metadata, - meta_file, - cache_dir, - IntegrationConfig.CIVITAI_API_KEY, - ) - - # Merge, local first, dedup - seen = set(preview_paths) - all_images = list(preview_paths) - for cp in civitai_paths: - if cp not in seen: - all_images.append(cp) - seen.add(cp) - - # Build prompt text: prefer example prompt, then model name - example_prompt = get_example_prompt_from_metadata(metadata) - prompt_text = example_prompt or model_name - - # Build tags - tags = ["lora-manager", f"lora:{model_name}"] - tags.extend(trigger_words) - - # Save to database via existing mechanism - try: - import hashlib - - prompt_hash = hashlib.sha256( - prompt_text.strip().lower().encode("utf-8") - ).hexdigest() - - existing = await self._run_in_executor( - self.db.get_prompt_by_hash, prompt_hash - ) - - if existing: - # Link all images - for pp in all_images: - await self._run_in_executor( - self.db.link_image_to_prompt, - existing["id"], - pp, - ) - skipped += 1 - else: - prompt_id = await self._run_in_executor( - self.db.save_prompt, - prompt_text, - "lora-manager", # category - tags, - None, # rating - None, # notes - prompt_hash, - ) - - if prompt_id: - for pp in all_images: - await self._run_in_executor( - self.db.link_image_to_prompt, - prompt_id, - pp, - ) - imported += 1 - else: - skipped += 1 - - except Exception as e: - self.logger.warning(f"Failed to import LoRA {model_name}: {e}") - skipped += 1 - - # Progress update for every LoRA - progress = int(5 + (90 * (i + 1) / max(total, 1))) - img_count = len(all_images) - status = f"{model_name}" - if img_count: - status += f" ({img_count} images)" - await send_progress( - { - "type": "progress", - "status": status, - "progress": progress, - "processed": i + 1, - "imported": imported, - "skipped": skipped, - } - ) - - await send_progress( - { - "type": "complete", - "progress": 100, - "total": total, - "imported": imported, - "skipped": skipped, - } - ) + try: + await self._lora_scan_stream(allowed, send_progress) + except Exception as e: + self.logger.exception("LoRA scan failed") + await send_progress({"type": "error", "message": self._public_error(e)}) await response.write_eof() return response except Exception as e: self.logger.error(f"LoRA scan failed: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return self._error_response(e) + + async def _lora_scan_stream(self, lm_path, send_progress): + """Import every LoRA under ``lm_path``, reporting through ``send_progress``.""" + from ..config import IntegrationConfig + from ..lora_utils import ( + download_civitai_images, + find_lora_directories, + get_example_prompt_from_metadata, + get_lora_image_cache_dir, + get_preview_images_from_metadata, + get_trigger_words_from_metadata, + get_model_name_from_metadata, + read_lora_metadata, + ) + + await send_progress( + { + "type": "progress", + "status": "Clearing previous lora-manager imports...", + "progress": 0, + } + ) + + # Clear previous imports so reimport is always clean + await self._run_in_executor(self.db.delete_prompts_by_category, "lora-manager") + + await send_progress( + { + "type": "progress", + "status": "Finding LoRA directories...", + "progress": 2, + } + ) + + lora_dirs = await self._run_in_executor(find_lora_directories, lm_path) + + # Collect all metadata files + meta_files = [] + for d in lora_dirs: + dir_path = Path(d) + meta_files.extend(dir_path.rglob("*.metadata.json")) + + total = len(meta_files) + imported = 0 + skipped = 0 + + await send_progress( + { + "type": "progress", + "status": f"Found {total} LoRA metadata files", + "progress": 5, + "total": total, + } + ) + + cache_dir = get_lora_image_cache_dir() + + for i, meta_file in enumerate(meta_files): + metadata = await self._run_in_executor(read_lora_metadata, meta_file) + if not metadata: + skipped += 1 + continue + + model_name = get_model_name_from_metadata(metadata) + trigger_words = get_trigger_words_from_metadata(metadata) + + # Collect all images: local previews + downloaded civitai examples + preview_paths = await self._run_in_executor( + get_preview_images_from_metadata, metadata, meta_file + ) + civitai_paths = await self._run_in_executor( + download_civitai_images, + metadata, + meta_file, + cache_dir, + IntegrationConfig.CIVITAI_API_KEY, + ) + + # Merge, local first, dedup + seen = set(preview_paths) + all_images = list(preview_paths) + for cp in civitai_paths: + if cp not in seen: + all_images.append(cp) + seen.add(cp) + + # Build prompt text: prefer example prompt, then model name + example_prompt = get_example_prompt_from_metadata(metadata) + prompt_text = example_prompt or model_name + + # Build tags + tags = ["lora-manager", f"lora:{model_name}"] + tags.extend(trigger_words) + + # Save to database via existing mechanism + try: + import hashlib + + prompt_hash = hashlib.sha256( + prompt_text.strip().lower().encode("utf-8") + ).hexdigest() + + existing = await self._run_in_executor( + self.db.get_prompt_by_hash, prompt_hash + ) + + if existing: + # Link all images + for pp in all_images: + await self._run_in_executor( + self.db.link_image_to_prompt, + existing["id"], + pp, + ) + skipped += 1 + else: + prompt_id = await self._run_in_executor( + self.db.save_prompt, + prompt_text, + "lora-manager", # category + tags, + None, # rating + None, # notes + prompt_hash, + ) + + if prompt_id: + for pp in all_images: + await self._run_in_executor( + self.db.link_image_to_prompt, + prompt_id, + pp, + ) + imported += 1 + else: + skipped += 1 + + except Exception as e: + self.logger.warning(f"Failed to import LoRA {model_name}: {e}") + skipped += 1 + + # Progress update for every LoRA + progress = int(5 + (90 * (i + 1) / max(total, 1))) + img_count = len(all_images) + status = f"{model_name}" + if img_count: + status += f" ({img_count} images)" + await send_progress( + { + "type": "progress", + "status": status, + "progress": progress, + "processed": i + 1, + "imported": imported, + "skipped": skipped, + } + ) + + await send_progress( + { + "type": "complete", + "progress": 100, + "total": total, + "imported": imported, + "skipped": skipped, + } + ) # ── Trigger word endpoints ─────────────────────────────────────── @@ -364,17 +413,11 @@ class LoraIntegrationMixin: from ..lora_utils import get_trigger_cache if not IntegrationConfig.LORA_MANAGER_ENABLED: - return web.json_response( - {"success": False, "error": "LoraManager integration not enabled"}, - status=400, - ) + return self._bad_lora_request("LoraManager integration not enabled") lora_name = request.query.get("name", "") if not lora_name: - return web.json_response( - {"success": False, "error": "Missing 'name' query parameter"}, - status=400, - ) + return self._bad_lora_request("Missing 'name' query parameter") cache = get_trigger_cache() if not cache.is_loaded: @@ -388,7 +431,7 @@ class LoraIntegrationMixin: ) except Exception as e: self.logger.error(f"Trigger word lookup failed: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return self._error_response(e) async def lora_refresh_cache(self, request): """Force-refresh the trigger word cache from disk.""" @@ -397,17 +440,11 @@ class LoraIntegrationMixin: from ..lora_utils import get_trigger_cache if not IntegrationConfig.LORA_MANAGER_ENABLED: - return web.json_response( - {"success": False, "error": "LoraManager integration not enabled"}, - status=400, - ) + return self._bad_lora_request("LoraManager integration not enabled") lm_path = IntegrationConfig.LORA_MANAGER_PATH if not lm_path: - return web.json_response( - {"success": False, "error": "LoraManager path not configured"}, - status=400, - ) + return self._bad_lora_request("LoraManager path not configured") cache = get_trigger_cache() count = await self._run_in_executor(cache.load, lm_path) @@ -416,4 +453,4 @@ class LoraIntegrationMixin: ) except Exception as e: self.logger.error(f"Trigger cache refresh failed: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return self._error_response(e) diff --git a/py/api/prompts.py b/py/api/prompts.py index 5d0efdf..4f4b401 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -2,8 +2,92 @@ import datetime import json +import os from aiohttp import web +from urllib.parse import quote + +# Upper bounds on page sizes and offsets for every list endpoint. The offset +# ceiling keeps values inside SQLite's 64-bit range and bounds skip work. +MAX_PAGE_LIMIT = 500 +MAX_PAGE_OFFSET = 10_000_000 + +# The export walks the whole table in pages of this size instead of asking +# for one giant result set (and silently truncating large libraries). +EXPORT_PAGE_SIZE = 1000 + + +def parse_page_params(query, default_limit=50, max_limit=MAX_PAGE_LIMIT): + """Return ``(limit, offset)`` from a query mapping, clamped to safe bounds. + + ``limit`` is clamped to ``[1, max_limit]`` and ``offset`` to + ``[0, MAX_PAGE_OFFSET]``. Non-integer values raise ``ValueError`` so the + caller can answer 400. + """ + limit = int(query.get("limit", default_limit)) + offset = int(query.get("offset", 0)) + return max(1, min(limit, max_limit)), max(0, min(offset, MAX_PAGE_OFFSET)) + + +def bad_request(message): + return web.json_response({"success": False, "error": message}, status=400) + + +def tags_error(tags): + """A 400 response when *tags* fail ``validate_tags``, else None.""" + try: + validate_tags(tags) + except ValueError as exc: + return bad_request(str(exc)) + return None + + +def safe_error_message(exc): + """Describe *exc* for a client without leaking absolute server paths. + + OSError carries the offending path in ``filename``; only its basename + is echoed back together with ``strerror``. + """ + if isinstance(exc, OSError): + reason = exc.strerror or exc.__class__.__name__ + if exc.filename: + return f"{reason}: {os.path.basename(str(exc.filename))}" + return reason + return str(exc) or exc.__class__.__name__ + + +def thumbnail_url_for(thumb_rel, thumb_abs): + """Serve URL for an existing thumbnail, or None when the file is missing. + + The URL carries the file's mtime as ``?v=``: thumbnails are served with a + long cache lifetime, so a regenerated file must get a new URL or browsers + keep showing whatever they cached for the old one. + """ + try: + version = int(thumb_abs.stat().st_mtime) + except OSError: + return None + return f"/prompt_manager/images/serve/{quote(thumb_rel, safe='/')}?v={version}" + + +def publish_image_paths(images, public_path): + """Image dicts for a response: ``image_path`` in its public (ComfyUI + relative) form and the absolute ``file_path`` duplicate removed. + + ``url``/``thumbnail_url`` added by enrichment are left untouched, so the + frontend keeps working without ever seeing the server's directory layout. + """ + published = [] + for image in images: + if not isinstance(image, dict): + published.append(image) + continue + image = {k: v for k, v in image.items() if k != "file_path"} + if "image_path" in image: + image["image_path"] = public_path(image["image_path"]) + published.append(image) + return published + try: from ...utils.validators import ( @@ -50,6 +134,10 @@ class PromptRoutesMixin: async def get_tags_filter_route(request): return await self.get_tags_filter(request) + @routes.get("/prompt_manager/tags/suggest") + async def suggest_tags_route(request): + return await self.suggest_tags(request) + # Bulk tag operations (register BEFORE {tag_name} to avoid path param match) @routes.post("/prompt_manager/tags/merge") async def merge_tags_route(request): @@ -122,6 +210,15 @@ class PromptRoutesMixin: async def export_prompts_route(request): return await self.export_prompts(request) + def _present_prompts(self, prompts): + """Prompts ready for a response: image urls added, server paths hidden.""" + self._enrich_prompt_images(prompts) + for prompt in prompts: + prompt["images"] = publish_image_paths( + prompt.get("images", []), self._public_path + ) + return prompts + async def search_prompts(self, request): """Search for prompts using multiple filter criteria.""" try: @@ -129,7 +226,10 @@ class PromptRoutesMixin: category = request.query.get("category", "").strip() tags_str = request.query.get("tags", "").strip() min_rating = request.query.get("min_rating", 0) - limit = int(request.query.get("limit", 50)) + try: + limit, offset = parse_page_params(request.query) + except ValueError: + return bad_request("limit and offset must be integers") folder = request.query.get("folder", "").strip() or None # Validated against a whitelist in the database layer @@ -151,19 +251,33 @@ class PromptRoutesMixin: tags=tags, rating_min=min_rating, limit=limit, + offset=offset, folder=folder, sort=sort, ) - self._enrich_prompt_images(results) + self._present_prompts(results) return web.json_response( - {"success": True, "results": results, "count": len(results)} + { + "success": True, + "results": results, + "count": len(results), + "pagination": { + "limit": limit, + "offset": offset, + "count": len(results), + }, + } ) except Exception as e: self.logger.error(f"Search error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Search failed: {str(e)}", "results": []}, + { + "success": False, + "error": f"Search failed: {safe_error_message(e)}", + "results": [], + }, status=500, ) @@ -182,28 +296,29 @@ class PromptRoutesMixin: return web.json_response({"success": True, "subfolders": subfolders}) except Exception as e: self.logger.error(f"Subfolders error: {e}") - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": safe_error_message(e)}, status=500 + ) async def get_recent_prompts(self, request): - """Retrieve prompts with pagination and an optional sort (default newest first).""" + """Retrieve prompts with pagination and an optional sort (default newest + first).""" try: - limit = int(request.query.get("limit", 50)) - page = int(request.query.get("page", 1)) - offset = int(request.query.get("offset", 0)) + try: + limit, offset = parse_page_params(request.query) + page = int(request.query.get("page", 1)) + except ValueError: + return bad_request("limit, offset and page must be integers") if page > 1 and offset == 0: - offset = (page - 1) * limit - - if limit > 1000: - limit = 1000 - elif limit < 1: - limit = 1 + offset = min((page - 1) * limit, MAX_PAGE_OFFSET) + # Validated against a whitelist in the database layer sort = request.query.get("sort") or None results = await self._run_in_executor( self.db.get_recent_prompts, limit=limit, offset=offset, sort=sort ) - self._enrich_prompt_images(results["prompts"]) + self._present_prompts(results["prompts"]) return web.json_response( { @@ -226,7 +341,7 @@ class PromptRoutesMixin: return web.json_response( { "success": False, - "error": f"Failed to get recent prompts: {str(e)}", + "error": f"Failed to get recent prompts: {safe_error_message(e)}", "results": [], "pagination": {"total": 0, "page": 1, "total_pages": 0}, }, @@ -243,7 +358,7 @@ class PromptRoutesMixin: return web.json_response( { "success": False, - "error": f"Failed to get categories: {str(e)}", + "error": f"Failed to get categories: {safe_error_message(e)}", "categories": [], }, status=500, @@ -259,7 +374,7 @@ class PromptRoutesMixin: return web.json_response( { "success": False, - "error": f"Failed to get tags: {str(e)}", + "error": f"Failed to get tags: {safe_error_message(e)}", "tags": [], }, status=500, @@ -269,13 +384,9 @@ class PromptRoutesMixin: """Get tags with usage counts, search, sort, and pagination.""" try: try: - limit = int(request.query.get("limit", 50)) - offset = int(request.query.get("offset", 0)) - except (ValueError, TypeError): - return web.json_response( - {"success": False, "error": "Invalid limit or offset parameter"}, - status=400, - ) + limit, offset = parse_page_params(request.query) + except ValueError: + return bad_request("limit and offset must be integers") search = request.query.get("search", "").strip() or None sort = request.query.get("sort", "alpha_asc") @@ -301,32 +412,30 @@ class PromptRoutesMixin: ) except Exception as e: self.logger.error(f"Tags stats error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": safe_error_message(e)}, status=500 + ) async def get_tag_prompts(self, request): """Get prompts for a single tag.""" try: - from urllib.parse import unquote - - tag_name = unquote(request.match_info.get("tag_name", "")) + # aiohttp has already percent-decoded match_info; decoding again + # would turn a tag literally named "%41" into "A". + tag_name = request.match_info.get("tag_name", "") if not tag_name: return web.json_response( {"success": False, "error": "Tag name required"}, status=400 ) try: - limit = int(request.query.get("limit", 20)) - offset = int(request.query.get("offset", 0)) - except (ValueError, TypeError): - return web.json_response( - {"success": False, "error": "Invalid limit or offset parameter"}, - status=400, - ) + limit, offset = parse_page_params(request.query, default_limit=20) + except ValueError: + return bad_request("limit and offset must be integers") result = await self._run_in_executor( self.db.get_prompts_by_tags, [tag_name], "and", limit, offset ) - self._enrich_prompt_images(result["prompts"]) + self._present_prompts(result["prompts"]) return web.json_response( { @@ -343,7 +452,46 @@ class PromptRoutesMixin: ) except Exception as e: self.logger.error(f"Tag prompts error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": safe_error_message(e)}, status=500 + ) + + async def suggest_tags(self, request): + """Autocomplete for the tag search box. + + ``q`` is the prefix being typed, ``with`` the comma-separated tags + already entered: only tags that co-occur with all of them are offered. + """ + try: + prefix = request.query.get("q", "").strip()[:100] + context = [] + seen = set() + for raw in request.query.get("with", "").split(","): + name = raw.strip() + if name and name.lower() not in seen: + seen.add(name.lower()) + context.append(name) + context = context[:10] + try: + limit = int(request.query.get("limit", 15)) + except ValueError: + limit = 15 + limit = max(1, min(limit, 50)) + + suggestions = await self._run_in_executor( + self.db.suggest_tags, prefix, context, limit + ) + return web.json_response({"success": True, "suggestions": suggestions}) + except Exception as e: + self.logger.error(f"Tag suggest error: {e}") + return web.json_response( + { + "success": False, + "error": "Failed to suggest tags", + "suggestions": [], + }, + status=500, + ) async def get_tags_filter(self, request): """Get prompts matching multiple tags with AND/OR mode, or untagged prompts.""" @@ -352,20 +500,13 @@ class PromptRoutesMixin: if untagged: try: - limit = int(request.query.get("limit", 20)) - offset = int(request.query.get("offset", 0)) - except (ValueError, TypeError): - return web.json_response( - { - "success": False, - "error": "Invalid limit or offset parameter", - }, - status=400, - ) + limit, offset = parse_page_params(request.query, default_limit=20) + except ValueError: + return bad_request("limit and offset must be integers") result = await self._run_in_executor( self.db.get_untagged_prompts, limit, offset ) - self._enrich_prompt_images(result["prompts"]) + self._present_prompts(result["prompts"]) return web.json_response( { "success": True, @@ -393,18 +534,14 @@ class PromptRoutesMixin: mode = "and" try: - limit = int(request.query.get("limit", 20)) - offset = int(request.query.get("offset", 0)) - except (ValueError, TypeError): - return web.json_response( - {"success": False, "error": "Invalid limit or offset parameter"}, - status=400, - ) + limit, offset = parse_page_params(request.query, default_limit=20) + except ValueError: + return bad_request("limit and offset must be integers") result = await self._run_in_executor( self.db.get_prompts_by_tags, tags_list, mode, limit, offset ) - self._enrich_prompt_images(result["prompts"]) + self._present_prompts(result["prompts"]) return web.json_response( { @@ -422,14 +559,16 @@ class PromptRoutesMixin: ) except Exception as e: self.logger.error(f"Tags filter error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": safe_error_message(e)}, status=500 + ) async def rename_tag_endpoint(self, request): """Rename a tag across all prompts.""" try: - from urllib.parse import unquote - - tag_name = unquote(request.match_info.get("tag_name", "")) + # aiohttp has already percent-decoded match_info; decoding again + # would turn a tag literally named "%41" into "A". + tag_name = request.match_info.get("tag_name", "") if not tag_name: return web.json_response( {"success": False, "error": "Tag name required"}, status=400 @@ -441,11 +580,15 @@ class PromptRoutesMixin: return web.json_response( {"success": False, "error": "Invalid JSON body"}, status=400 ) - new_name = (body.get("new_name") or "").strip() - if not new_name: + raw_name = body.get("new_name") + if not raw_name: return web.json_response( {"success": False, "error": "New tag name required"}, status=400 ) + error = tags_error([raw_name]) + if error is not None: + return error + new_name = raw_name.strip() result = await self._run_in_executor( self.db.rename_tag_all_prompts, tag_name, new_name @@ -459,19 +602,22 @@ class PromptRoutesMixin: if result.get("skipped_count", 0) > 0: resp["skipped_count"] = result["skipped_count"] resp["warning"] = ( - f"{result['skipped_count']} prompt(s) had corrupted tag data and were skipped" + f"{result['skipped_count']} prompt(s) had corrupted tag data " + "and were skipped" ) return web.json_response(resp) except Exception as e: self.logger.error(f"Rename tag error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": safe_error_message(e)}, status=500 + ) async def delete_tag_endpoint(self, request): """Delete a tag from all prompts.""" try: - from urllib.parse import unquote - - tag_name = unquote(request.match_info.get("tag_name", "")) + # aiohttp has already percent-decoded match_info; decoding again + # would turn a tag literally named "%41" into "A". + tag_name = request.match_info.get("tag_name", "") if not tag_name: return web.json_response( {"success": False, "error": "Tag name required"}, status=400 @@ -488,12 +634,15 @@ class PromptRoutesMixin: if result.get("skipped_count", 0) > 0: resp["skipped_count"] = result["skipped_count"] resp["warning"] = ( - f"{result['skipped_count']} prompt(s) had corrupted tag data and were skipped" + f"{result['skipped_count']} prompt(s) had corrupted tag data " + "and were skipped" ) return web.json_response(resp) except Exception as e: self.logger.error(f"Delete tag error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": safe_error_message(e)}, status=500 + ) async def merge_tags_endpoint(self, request): """Merge source tags into a target tag.""" @@ -505,16 +654,22 @@ class PromptRoutesMixin: {"success": False, "error": "Invalid JSON body"}, status=400 ) source_tags = body.get("source_tags", []) - target_tag = (body.get("target_tag") or "").strip() + raw_target = body.get("target_tag") if not source_tags: return web.json_response( {"success": False, "error": "Source tags required"}, status=400 ) - if not target_tag: + if not raw_target: return web.json_response( {"success": False, "error": "Target tag required"}, status=400 ) + if not isinstance(source_tags, list): + return bad_request("Source tags must be a list") + error = tags_error([*source_tags, raw_target]) + if error is not None: + return error + target_tag = raw_target.strip() result = await self._run_in_executor( self.db.merge_tags, source_tags, target_tag @@ -528,12 +683,15 @@ class PromptRoutesMixin: if result.get("skipped_count", 0) > 0: resp["skipped_count"] = result["skipped_count"] resp["warning"] = ( - f"{result['skipped_count']} prompt(s) had corrupted tag data and were skipped" + f"{result['skipped_count']} prompt(s) had corrupted tag data " + "and were skipped" ) return web.json_response(resp) except Exception as e: self.logger.error(f"Merge tags error: {e}", exc_info=True) - return web.json_response({"success": False, "error": str(e)}, status=500) + return web.json_response( + {"success": False, "error": safe_error_message(e)}, status=500 + ) async def save_prompt(self, request): """Save a new prompt with metadata and duplicate detection.""" @@ -608,7 +766,10 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Save error: {e}", exc_info=True) return web.json_response( - {"success": False, "error": f"Failed to save prompt: {str(e)}"}, + { + "success": False, + "error": f"Failed to save prompt: {safe_error_message(e)}", + }, status=500, ) @@ -638,7 +799,10 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Delete error: {e}") return web.json_response( - {"success": False, "error": f"Failed to delete prompt: {str(e)}"}, + { + "success": False, + "error": f"Failed to delete prompt: {safe_error_message(e)}", + }, status=500, ) @@ -682,7 +846,10 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Update prompt error: {e}") return web.json_response( - {"success": False, "error": f"Failed to update prompt: {str(e)}"}, + { + "success": False, + "error": f"Failed to update prompt: {safe_error_message(e)}", + }, status=500, ) @@ -719,7 +886,10 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Update rating error: {e}") return web.json_response( - {"success": False, "error": f"Failed to update rating: {str(e)}"}, + { + "success": False, + "error": f"Failed to update rating: {safe_error_message(e)}", + }, status=500, ) @@ -728,12 +898,16 @@ class PromptRoutesMixin: try: prompt_id = int(request.match_info["prompt_id"]) data = await request.json() - new_tag = (data.get("tag") or "").strip() + raw_tag = data.get("tag") - if not new_tag: + if not raw_tag: return web.json_response( {"success": False, "error": "Tag cannot be empty"}, status=400 ) + error = tags_error([raw_tag]) + if error is not None: + return error + new_tag = raw_tag.strip() prompt = await self._run_in_executor(self.db.get_prompt_by_id, prompt_id) if not prompt: @@ -762,7 +936,11 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Add tag error: {e}") return web.json_response( - {"success": False, "error": f"Failed to add tag: {str(e)}"}, status=500 + { + "success": False, + "error": f"Failed to add tag: {safe_error_message(e)}", + }, + status=500, ) async def add_tags_to_prompt(self, request): @@ -782,6 +960,12 @@ class PromptRoutesMixin: {"success": False, "error": "Tags must be a non-empty list"}, status=400, ) + # Blank entries are skipped below, so only the rest is validated. + error = tags_error( + [t for t in new_tags if not (isinstance(t, str) and not t.strip())] + ) + if error is not None: + return error prompt = await self._run_in_executor(self.db.get_prompt_by_id, prompt_id) if not prompt: @@ -820,7 +1004,11 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Add tags error: {e}") return web.json_response( - {"success": False, "error": f"Failed to add tags: {str(e)}"}, status=500 + { + "success": False, + "error": f"Failed to add tags: {safe_error_message(e)}", + }, + status=500, ) async def remove_prompt_tag(self, request): @@ -857,7 +1045,10 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Remove tag error: {e}") return web.json_response( - {"success": False, "error": f"Failed to remove tag: {str(e)}"}, + { + "success": False, + "error": f"Failed to remove tag: {safe_error_message(e)}", + }, status=500, ) @@ -887,7 +1078,10 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Bulk delete error: {e}") return web.json_response( - {"success": False, "error": f"Failed to delete prompts: {str(e)}"}, + { + "success": False, + "error": f"Failed to delete prompts: {safe_error_message(e)}", + }, status=500, ) @@ -903,6 +1097,11 @@ class PromptRoutesMixin: {"success": False, "error": "No prompt IDs or tags provided"}, status=400, ) + if not isinstance(new_tags, list): + return bad_request("Tags must be a list") + error = tags_error(new_tags) + if error is not None: + return error updated_count = await self._run_in_executor( self.db.bulk_add_tags, prompt_ids, new_tags @@ -919,7 +1118,11 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Bulk add tags error: {e}") return web.json_response( - {"success": False, "error": f"Failed to add tags: {str(e)}"}, status=500 + { + "success": False, + "error": f"Failed to add tags: {safe_error_message(e)}", + }, + status=500, ) async def bulk_set_category(self, request): @@ -949,14 +1152,27 @@ class PromptRoutesMixin: except Exception as e: self.logger.error(f"Bulk set category error: {e}") return web.json_response( - {"success": False, "error": f"Failed to set category: {str(e)}"}, + { + "success": False, + "error": f"Failed to set category: {safe_error_message(e)}", + }, status=500, ) + def _all_prompts_for_export(self): + """Every prompt, fetched in EXPORT_PAGE_SIZE pages (blocking).""" + prompts, offset = [], 0 + while True: + page = self.db.search_prompts(limit=EXPORT_PAGE_SIZE, offset=offset) + prompts.extend(page) + if len(page) < EXPORT_PAGE_SIZE: + return prompts + offset += EXPORT_PAGE_SIZE + async def export_prompts(self, request): """Export all prompts to JSON.""" try: - prompts = await self._run_in_executor(self.db.search_prompts, limit=10000) + prompts = await self._run_in_executor(self._all_prompts_for_export) export_data = { "export_date": datetime.datetime.now(datetime.timezone.utc).isoformat(), @@ -970,13 +1186,19 @@ class PromptRoutesMixin: text=json_data, content_type="application/json", headers={ - "Content-Disposition": f'attachment; filename="prompt_manager_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}.json"' + "Content-Disposition": ( + 'attachment; filename="prompt_manager_' + f'{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}.json"' + ) }, ) except Exception as e: self.logger.error(f"Export error: {e}") return web.json_response( - {"success": False, "error": f"Failed to export prompts: {str(e)}"}, + { + "success": False, + "error": f"Failed to export prompts: {safe_error_message(e)}", + }, status=500, ) diff --git a/py/api/scan_job.py b/py/api/scan_job.py new file mode 100644 index 0000000..d3eb9ce --- /dev/null +++ b/py/api/scan_job.py @@ -0,0 +1,61 @@ +"""Progress fan-out for the background output scan. + +The scan runs as an asyncio task that outlives the HTTP request that started +it, so an accidental tab close does not abandon a half-finished scan. Every +SSE subscriber (the request that started the scan, a reconnecting page, a +second tab) receives the latest event as a snapshot and then every event +published after it, until the job finishes. +""" + +import asyncio + + +class ScanJob: + """One scan run: its latest event, completion flag and live subscribers.""" + + def __init__(self): + self.last_event = None + self.finished = False + self.task = None + self._subscribers = set() + + @property + def running(self): + return not self.finished + + @property + def subscriber_count(self): + return len(self._subscribers) + + def publish(self, payload): + """Record ``payload`` as the latest event and hand it to every subscriber.""" + self.last_event = payload + for queue in list(self._subscribers): + queue.put_nowait(payload) + + def finish(self): + """Mark the job done and release every subscriber after its queued events.""" + self.finished = True + for queue in list(self._subscribers): + queue.put_nowait(None) + + async def events(self): + """Yield the latest event, then each later one, until the job finishes. + + Consume it under ``contextlib.aclosing`` so a consumer that stops early + (a closed socket) drops its queue at once instead of at garbage collection. + """ + queue = asyncio.Queue() + self._subscribers.add(queue) + try: + if self.last_event is not None: + yield self.last_event + if self.finished: + return + while True: + payload = await queue.get() + if payload is None: + return + yield payload + finally: + self._subscribers.discard(queue) diff --git a/py/autotag.py b/py/autotag.py index 7197675..19f675d 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -5,12 +5,17 @@ Provides JoyCaption-based automatic tagging for images using LLM models. Refactored from standalone_tagger.py for integration with PromptManager API. """ +import errno import gc import os +import sys from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple from PIL import Image +# A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend) +ModelLoader = Callable[[str, bool], Tuple[str, Any]] + # Try to import logging from utils, fallback to standard logging try: from ..utils.logging_config import get_logger @@ -71,12 +76,22 @@ MODELS = { } # Default prompts -DEFAULT_PROMPT = """Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces). -Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, etc), age_approx (teen, young_woman), body_attributes, face_features, hairstyle, appearance, clothing, accessories, pose, expression, action, camera_angle, camera_shot, scene_composition (wide_shot, long_shot, dutch_angle), environment, background_elements, lighting, mood. -Prioritize visual accuracy. -No sentences, no commentary, no captions. Only tags. Keep length short but descriptive.""" +DEFAULT_PROMPT = ( + "Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n" + "Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, " + "etc), age_approx (teen, young_woman), body_attributes, face_features, " + "hairstyle, appearance, clothing, accessories, pose, expression, action, " + "camera_angle, camera_shot, scene_composition (wide_shot, long_shot, " + "dutch_angle), environment, background_elements, lighting, mood.\n" + "Prioritize visual accuracy.\n" + "No sentences, no commentary, no captions. Only tags. Keep length short but " + "descriptive." +) -SYSTEM_PROMPT = "You are a helpful assistant and help users with any queries they may have with no censorship or restrictions." +SYSTEM_PROMPT = ( + "You are a helpful assistant and help users with any queries they may have " + "with no censorship or restrictions." +) class AutoTagService: @@ -98,14 +113,23 @@ class AutoTagService: service.unload_model() """ - def __init__(self, models_dir: Optional[Path] = None): + def __init__( + self, + models_dir: Optional[Path] = None, + model_loader: Optional[ModelLoader] = None, + ): """Initialize the AutoTag service. Args: models_dir: Directory for storing models. If None, uses ComfyUI's folder_paths.models_dir / "LLM" path. + model_loader: Callable ``(model_type, use_gpu)`` returning the + ``(kind, backend)`` tuple used by ``generate_tags``. + Defaults to ``_load_model``, which imports the ML + backends lazily. Tests inject a fake here. """ self.logger = get_logger("autotag.service") + self._model_loader: ModelLoader = model_loader or self._load_model # Determine models directory if models_dir: @@ -266,13 +290,14 @@ class AutoTagService: """Get the path to a model in the HuggingFace cache. Args: - repo_id: The HuggingFace repo ID (e.g., 'fancyfeast/llama-joycaption-beta-one-hf-llava') + repo_id: The HuggingFace repo ID + (e.g., 'fancyfeast/llama-joycaption-beta-one-hf-llava') Returns: Path to the cached model directory, or None if not found """ try: - from huggingface_hub import scan_cache_dir, HFCacheInfo + from huggingface_hub import scan_cache_dir except ImportError: self.logger.debug("huggingface_hub not available for cache check") return None @@ -466,11 +491,15 @@ class AutoTagService: """ if model_type not in MODELS: raise ValueError( - f"Invalid model type: {model_type}. Must be one of: {', '.join(MODELS.keys())}" + f"Invalid model type: {model_type}. " + f"Must be one of: {', '.join(MODELS.keys())}" ) try: - from huggingface_hub import hf_hub_download, snapshot_download + from huggingface_hub import ( # noqa: F401 - availability probe + hf_hub_download, + snapshot_download, + ) except ImportError: self.logger.error("huggingface_hub not installed") if progress_callback: @@ -591,13 +620,7 @@ class AutoTagService: raise RuntimeError(f"Model {model_type} not downloaded") try: - if model_type == "gguf": - self._tagger = self._load_gguf_tagger(use_gpu) - elif model_type.startswith("wd14"): - self._tagger = self._load_wd14_tagger(model_type, use_gpu) - else: - self._tagger = self._load_hf_tagger() - + self._tagger = self._model_loader(model_type, use_gpu) self._current_model_type = model_type self.logger.info(f"Model {model_type} loaded successfully") return True @@ -608,6 +631,14 @@ class AutoTagService: self._current_model_type = None raise RuntimeError(f"Failed to load model: {e}") + def _load_model(self, model_type: str, use_gpu: bool = True) -> Tuple[str, Any]: + """Default loader: dispatch to a backend, importing it only now.""" + if model_type == "gguf": + return self._load_gguf_tagger(use_gpu) + if model_type.startswith("wd14"): + return self._load_wd14_tagger(model_type, use_gpu) + return self._load_hf_tagger() + def _load_gguf_tagger(self, use_gpu: bool = True): """Load GGUF-based tagger.""" from llama_cpp import Llama @@ -690,14 +721,11 @@ class AutoTagService: self._current_model_type = None gc.collect() - # Try to clear CUDA cache if available - try: - import torch - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - except ImportError: - pass + # Clear the CUDA cache only if torch is already loaded; importing + # it here would pull in a heavy dependency the backend never used. + torch = sys.modules.get("torch") + if torch is not None and torch.cuda.is_available(): + torch.cuda.empty_cache() self.logger.info("Model unloaded, memory freed") @@ -720,7 +748,8 @@ class AutoTagService: Args: image_path: Path to the image file - prompt: Custom prompt for tag generation (LLM models only). Uses default if None. + prompt: Custom prompt for tag generation (LLM models only). + Uses default if None. general_threshold: Confidence threshold for general tags (WD14 only). character_threshold: Confidence threshold for character tags (WD14 only). @@ -735,7 +764,7 @@ class AutoTagService: raise RuntimeError("No model loaded. Call load_model() first.") if not os.path.exists(image_path): - raise FileNotFoundError(f"Image not found: {image_path}") + raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path)) use_prompt = prompt or self._custom_prompt diff --git a/py/config.py b/py/config.py index 88ef6b9..cc89108 100644 --- a/py/config.py +++ b/py/config.py @@ -22,7 +22,7 @@ Example: extension_name = "PromptManager" # Get server instance and routes (same pattern as ComfyUI_Assets) -from server import PromptServer +from server import PromptServer # noqa: E402 server_instance = PromptServer.instance routes = server_instance.routes @@ -30,8 +30,9 @@ routes = server_instance.routes # Extension info extension_uri = None # Will be set in __init__.py -import os -from typing import Dict, Any, List +import os # noqa: E402 +from pathlib import Path # noqa: E402 +from typing import Any, Dict, List, Optional, Tuple # noqa: E402 # Import logging system try: @@ -46,6 +47,108 @@ except ImportError: # Initialize logger for config operations config_logger = get_logger("prompt_manager.config") +# Environment variable listing extra directories (os.pathsep-separated) that +# may be used as gallery roots in addition to ComfyUI's own directories. +EXTRA_GALLERY_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS" + +# Environment variable overriding where config.json is read from and saved to. +CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH" + +# ComfyUI folder_paths getters whose directories are valid gallery parents. +_COMFYUI_DIRECTORY_GETTERS = ( + "get_output_directory", + "get_input_directory", + "get_temp_directory", + "get_user_directory", +) + + +def _canonical_path(path: str) -> str: + """Return a path normalised for comparison (realpath + normcase).""" + return os.path.normcase(os.path.realpath(path)) + + +def write_private_json(path: str, data: Any) -> None: + """Write ``data`` as JSON to ``path`` readable by the current user only. + + config.json can hold an API key (integrations.lora_manager), so it is + created with mode 0600 and, on POSIX, an existing more open file is + tightened too. Windows has no POSIX modes; there the chmod is a best + effort that only affects the read-only bit. Parent directories are + created as needed. Raises OSError when the file cannot be written. + """ + import json + + parent_dir = os.path.dirname(path) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2) + if os.name == "posix": + os.chmod(path, 0o600) + + +def _is_filesystem_root(path: str) -> bool: + """True for '/' on POSIX and drive roots such as 'C:\\' on Windows.""" + return os.path.dirname(path) == path + + +def _import_folder_paths(): + """ComfyUI's folder_paths module, or None outside ComfyUI.""" + try: + import folder_paths + except ImportError: + return None + return folder_paths + + +def _comfyui_base_directory() -> Optional[str]: + """ComfyUI's base directory: folder_paths.base_path, else the output dir's + parent.""" + folder_paths = _import_folder_paths() + if folder_paths is None: + return None + base = getattr(folder_paths, "base_path", None) + if isinstance(base, str) and base: + return base + getter = getattr(folder_paths, "get_output_directory", None) + if getter is None: + return None + try: + output_dir = getter() + except Exception: + return None + if isinstance(output_dir, str) and output_dir: + return os.path.dirname(os.path.realpath(output_dir)) + return None + + +def _comfyui_directories() -> List[str]: + """Directories reported by ComfyUI's folder_paths module, if importable.""" + folder_paths = _import_folder_paths() + if folder_paths is None: + return [] + + found = [] + for getter_name in _COMFYUI_DIRECTORY_GETTERS: + getter = getattr(folder_paths, getter_name, None) + if getter is None: + continue + try: + directory = getter() + except Exception: + continue + if isinstance(directory, str) and directory: + found.append(directory) + return found + + +def _extra_gallery_roots() -> List[str]: + """Directories the user opted in through PROMPT_MANAGER_EXTRA_GALLERY_ROOTS.""" + raw = os.environ.get(EXTRA_GALLERY_ROOTS_ENV, "") + return [entry.strip() for entry in raw.split(os.pathsep) if entry.strip()] + class GalleryConfig: """Configuration class for the gallery monitoring and image processing system. @@ -106,8 +209,9 @@ class GalleryConfig: Returns: Dict[str, Any]: A nested dictionary containing all gallery configuration - sections: monitoring, tracking, database, web_interface, and performance. - Each section contains the relevant configuration parameters as key-value pairs. + sections: monitoring, tracking, database, web_interface, and + performance. Each section contains the relevant configuration + parameters as key-value pairs. Example: config = GalleryConfig.get_config() @@ -166,7 +270,9 @@ class GalleryConfig: if "enabled" in monitoring: cls.MONITORING_ENABLED = monitoring["enabled"] if "directories" in monitoring: - cls.MONITORING_DIRECTORIES = monitoring["directories"] + directories = cls._validated_directories(monitoring["directories"]) + if directories is not None: + cls.MONITORING_DIRECTORIES = directories if "extensions" in monitoring: cls.SUPPORTED_EXTENSIONS = monitoring["extensions"] if "processing_delay" in monitoring: @@ -200,6 +306,127 @@ class GalleryConfig: if "metadata_extraction_timeout" in performance: cls.METADATA_EXTRACTION_TIMEOUT = performance["metadata_extraction_timeout"] + @classmethod + def _validated_directories(cls, directories: Any) -> Optional[List[str]]: + """Gallery roots from a config payload that pass validate_gallery_root. + + A hand-edited config.json is as untrusted as a settings request: + entries that fail validation are dropped with a warning (naming only + the basename), and a value that is not a list is ignored entirely + (``None`` is returned so the current roots stay as they are). + """ + if not isinstance(directories, list): + config_logger.warning( + "Ignoring gallery monitoring directories: expected a list" + ) + return None + kept = [] + for entry in directories: + ok, reason = cls.validate_gallery_root(entry) + if ok: + kept.append(entry) + continue + label = os.path.basename(entry) if isinstance(entry, str) else entry + config_logger.warning(f"Ignoring gallery root {label!r}: {reason}") + return kept + + @classmethod + def allowed_gallery_parents(cls) -> List[str]: + """Canonical directories under which gallery roots may live. + + Filesystem roots are never allowed as parents, even when listed in + the environment, because that would re-open the whole disk. + """ + parents = [] + for candidate in _comfyui_directories() + _extra_gallery_roots(): + try: + canonical = _canonical_path(candidate) + except (OSError, ValueError): + continue + if _is_filesystem_root(canonical) or not os.path.isdir(canonical): + continue + if canonical not in parents: + parents.append(canonical) + return parents + + @classmethod + def path_anchors(cls) -> List[str]: + """Canonical directories that relative gallery paths are resolved against. + + The ComfyUI base directory comes first, followed by the parent of each + directory listed in PROMPT_MANAGER_EXTRA_GALLERY_ROOTS, so an extra + root is addressed by its own name (``gallery/sub``) rather than by an + absolute path. The same anchors drive the public (relative) form of + paths in API responses. + """ + anchors = [] + base = _comfyui_base_directory() + candidates = [base] if base else [] + candidates.extend(os.path.dirname(root) for root in _extra_gallery_roots()) + for candidate in candidates: + try: + canonical = _canonical_path(candidate) + except (OSError, ValueError): + continue + if canonical not in anchors: + anchors.append(canonical) + return anchors + + @classmethod + def resolve_gallery_root(cls, path: str) -> str: + """Canonical absolute path for a gallery root given in any accepted form. + + Absolute paths are canonicalised as-is. Relative paths are tried + against each :meth:`path_anchors` entry and the first existing match + wins; otherwise the first anchor (or the current directory when there + is none) is used, so the caller's existence check reports it. + """ + path = path.strip() + if os.path.isabs(path): + return _canonical_path(path) + anchors = cls.path_anchors() + for anchor in anchors: + candidate = os.path.join(anchor, path) + if os.path.exists(candidate): + return _canonical_path(candidate) + return _canonical_path(os.path.join(anchors[0], path) if anchors else path) + + @classmethod + def validate_gallery_root(cls, path: Any) -> Tuple[bool, str]: + """Check whether ``path`` may be used as a gallery root. + + Returns: + (True, "") when the directory lies inside one of + :meth:`allowed_gallery_parents`; otherwise (False, reason). + """ + if not isinstance(path, str) or not path.strip(): + return False, "Gallery root must be a non-empty path" + + try: + canonical = cls.resolve_gallery_root(path) + except (OSError, ValueError): + return False, "Gallery root could not be resolved" + + if _is_filesystem_root(canonical): + return False, "Gallery root cannot be a filesystem root" + if canonical == _canonical_path(os.path.expanduser("~")): + return False, "Gallery root cannot be the home directory" + if not os.path.exists(canonical): + return False, "Gallery root does not exist" + if not os.path.isdir(canonical): + return False, "Gallery root is not a directory" + + candidate = Path(canonical) + for parent in cls.allowed_gallery_parents(): + if candidate.is_relative_to(Path(parent)): + return True, "" + + return ( + False, + "Gallery root must be inside a ComfyUI directory " + f"(output, input, temp, user) or one listed in {EXTRA_GALLERY_ROOTS_ENV}", + ) + class IntegrationConfig: """Configuration for third-party extension integrations. @@ -272,11 +499,36 @@ class PromptManagerConfig: RESULT_TIMEOUT = 5 # Seconds to auto-hide results in ComfyUI node SHOW_TEST_BUTTON = False # Show API test button in node UI WEBUI_DISPLAY_MODE = "newtab" # 'popup' or 'newtab' + INFINITE_SCROLL = False # Dashboard list: load the next page on scroll by default # Performance settings MAX_SEARCH_RESULTS = 100 ENABLE_FUZZY_SEARCH = False # Requires fuzzywuzzy AUTO_BACKUP_INTERVAL = 24 # Hours + # Threads for the output scan and thumbnail generation; half the cores by + # default so ComfyUI's own generation keeps CPU headroom. 1..max_worker_threads(). + DEFAULT_WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2) + WORKER_THREADS = DEFAULT_WORKER_THREADS + + @classmethod + def max_worker_threads(cls) -> int: + """Largest accepted worker-thread count: the detected core count.""" + return os.cpu_count() or 1 + + @classmethod + def clamp_worker_threads(cls, value, fallback: int) -> int: + """Coerce a worker count from config.json into ``[1, max_worker_threads()]``. + + Non-numeric values (and booleans) return ``fallback`` so a hand-edited + file cannot disable the setting. + """ + if isinstance(value, bool): + return fallback + if isinstance(value, str) and value.strip().isdigit(): + value = int(value.strip()) + if not isinstance(value, int): + return fallback + return max(1, min(value, cls.max_worker_threads())) @classmethod def get_config(cls) -> Dict[str, Any]: @@ -304,18 +556,33 @@ class PromptManagerConfig: "result_timeout": cls.RESULT_TIMEOUT, "show_test_button": cls.SHOW_TEST_BUTTON, "webui_display_mode": cls.WEBUI_DISPLAY_MODE, + "infinite_scroll": cls.INFINITE_SCROLL, }, "performance": { "max_search_results": cls.MAX_SEARCH_RESULTS, "enable_fuzzy_search": cls.ENABLE_FUZZY_SEARCH, "auto_backup_interval": cls.AUTO_BACKUP_INTERVAL, + "worker_threads": cls.WORKER_THREADS, }, "gallery": GalleryConfig.get_config(), "integrations": IntegrationConfig.get_config(), } @classmethod - def load_from_file(cls, config_path: str): + def get_config_path(cls) -> str: + """Path of the persisted config.json. + + Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file + lives at the repository root next to ``pyproject.toml``. + """ + override = os.environ.get(CONFIG_PATH_ENV, "").strip() + if override: + return override + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + return os.path.join(repo_root, "config.json") + + @classmethod + def load_from_file(cls, config_path: Optional[str] = None): """Load configuration settings from a JSON file. Reads configuration from the specified JSON file and updates the current @@ -324,7 +591,8 @@ class PromptManagerConfig: Args: config_path (str): Path to the JSON configuration file to load. - Can be relative or absolute path. + Can be relative or absolute path. Defaults to + :meth:`get_config_path`. Raises: The method handles all exceptions internally and logs errors rather @@ -336,6 +604,9 @@ class PromptManagerConfig: """ import json + if config_path is None: + config_path = cls.get_config_path() + if os.path.exists(config_path): try: with open(config_path, "r") as f: @@ -348,7 +619,7 @@ class PromptManagerConfig: config_logger.info(f"Config file not found: {config_path}, using defaults") @classmethod - def save_to_file(cls, config_path: str): + def save_to_file(cls, config_path: Optional[str] = None): """Save the current configuration to a JSON file. Serializes the complete configuration (including gallery settings) to @@ -357,6 +628,7 @@ class PromptManagerConfig: Args: config_path (str): Path where the JSON configuration file should be saved. Parent directories will be created if they don't exist. + Defaults to :meth:`get_config_path`. Raises: The method handles all exceptions internally and logs errors rather @@ -366,15 +638,11 @@ class PromptManagerConfig: PromptManagerConfig.save_to_file('backup_config.json') PromptManagerConfig.save_to_file('/etc/comfyui/prompt_manager.json') """ - import json + if config_path is None: + config_path = cls.get_config_path() try: - config = cls.get_config() - os.makedirs(os.path.dirname(config_path), exist_ok=True) - - with open(config_path, "w") as f: - json.dump(config, f, indent=2) - + write_private_json(config_path, cls.get_config()) config_logger.info(f"Saved configuration to {config_path}") except Exception as e: config_logger.error(f"Error saving config to {config_path}: {e}") @@ -390,7 +658,8 @@ class PromptManagerConfig: Args: new_config (Dict[str, Any]): Nested dictionary containing configuration updates. Should follow the same structure as returned by get_config(). - Valid top-level keys are: 'database', 'web_ui', 'performance', 'gallery'. + Valid top-level keys are: 'database', 'web_ui', 'performance', + 'gallery'. Example: new_settings = { @@ -414,6 +683,8 @@ class PromptManagerConfig: cls.SHOW_TEST_BUTTON = web_ui["show_test_button"] if "webui_display_mode" in web_ui: cls.WEBUI_DISPLAY_MODE = web_ui["webui_display_mode"] + if isinstance(web_ui.get("infinite_scroll"), bool): + cls.INFINITE_SCROLL = web_ui["infinite_scroll"] performance = new_config.get("performance", {}) if "max_search_results" in performance: @@ -422,6 +693,10 @@ class PromptManagerConfig: cls.ENABLE_FUZZY_SEARCH = performance["enable_fuzzy_search"] if "auto_backup_interval" in performance: cls.AUTO_BACKUP_INTERVAL = performance["auto_backup_interval"] + if "worker_threads" in performance: + cls.WORKER_THREADS = cls.clamp_worker_threads( + performance["worker_threads"], cls.WORKER_THREADS + ) # Update gallery config if "gallery" in new_config: @@ -434,8 +709,6 @@ class PromptManagerConfig: # Load configuration on import try: - config_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - config_file = os.path.join(config_dir, "config.json") - PromptManagerConfig.load_from_file(config_file) + PromptManagerConfig.load_from_file() except Exception as e: config_logger.error(f"Error during config initialization: {e}") diff --git a/py/lora_utils.py b/py/lora_utils.py index 62c925f..0c8b33e 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -15,6 +15,7 @@ import threading import urllib.request from pathlib import Path from typing import Dict, List, Optional, Tuple +from urllib.parse import urlsplit try: from ..utils.logging_config import get_logger @@ -108,6 +109,103 @@ def _looks_like_lora_manager(path: Path) -> bool: return (path / "py").is_dir() or (path / "lora_manager").is_dir() +def _public_path_anchors() -> List[str]: + """Directories the API's public (relative) paths are rendered against.""" + try: + from .config import GalleryConfig + except ImportError: + try: + from config import GalleryConfig + except ImportError: + return [] + try: + return list(GalleryConfig.path_anchors()) + except Exception: + return [] + + +def custom_nodes_directories() -> List[Path]: + """ComfyUI ``custom_nodes`` directories a LoraManager install may live in. + + Inside ComfyUI these come from ``folder_paths`` (``base_path`` plus any + registered ``custom_nodes`` folders); outside it, the parent of this + package, which is ``custom_nodes`` for a normal install. + """ + found: List[Path] = [] + + def add(candidate) -> None: + try: + resolved = Path(candidate).resolve() + except (OSError, RuntimeError, ValueError): + return + if resolved not in found: + found.append(resolved) + + try: + import folder_paths + except ImportError: + folder_paths = None + if folder_paths is not None: + base = getattr(folder_paths, "base_path", None) + if isinstance(base, str) and base: + add(Path(base) / "custom_nodes") + getter = getattr(folder_paths, "get_folder_paths", None) + if callable(getter): + try: + registered = getter("custom_nodes") or [] + except Exception: + registered = [] + for entry in registered: + add(entry) + if not found: + add(Path(__file__).resolve().parent.parent.parent) + return found + + +def resolve_lora_manager_path(path) -> Optional[str]: + """Canonical form of ``path`` when it may be used as the LoraManager install. + + Accepted: an existing directory that is a direct child of one of + :func:`custom_nodes_directories` and whose name contains ``lora`` + (case-insensitive). A relative path is tried against the ComfyUI root + and each custom_nodes directory. Anything else, including a symlink + that escapes custom_nodes, yields None. + """ + if not isinstance(path, str) or not path.strip(): + return None + path = path.strip() + roots = [] + for root in custom_nodes_directories(): + try: + roots.append(Path(os.path.normcase(os.path.realpath(str(root))))) + except (OSError, ValueError): + continue + + candidates = [path] + if not os.path.isabs(path): + # Public (API) paths are relative to these anchors; the ComfyUI root + # and the custom_nodes directories cover hand-typed relative paths. + bases = list(_public_path_anchors()) + comfy_root = find_comfyui_root() + if comfy_root: + bases.append(comfy_root) + bases.extend(root.parent for root in roots) + bases.extend(roots) + candidates = [os.path.join(str(base), path) for base in bases] + + for candidate in candidates: + try: + canonical = Path(os.path.normcase(os.path.realpath(candidate))) + is_dir = canonical.is_dir() + except (OSError, RuntimeError, ValueError): + continue + if not is_dir or "lora" not in canonical.name.lower(): + continue + if canonical.parent in roots: + return str(canonical) + return None + + # ── Metadata reading ───────────────────────────────────────────────── @@ -298,16 +396,95 @@ def get_preview_image_from_metadata( _THUMB_MAX_SIZE = 512 +# Only these hosts may ever receive the CivitAI API key, and only over HTTPS. +CIVITAI_HOSTS = frozenset( + {"civitai.com", "www.civitai.com", "api.civitai.com", "image.civitai.com"} +) +MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024 +_DOWNLOAD_CHUNK_BYTES = 64 * 1024 + + +class _RefuseRedirects(urllib.request.HTTPRedirectHandler): + """Never follow a redirect: urllib would re-send the Bearer header to it.""" + + def redirect_request(self, req, fp, code, msg, headers, newurl): + logger.warning(f"Refusing redirect from {req.full_url} to {newurl}") + return None + + +def _build_opener(*handlers) -> urllib.request.OpenerDirector: + """Build the outbound opener; extra handlers let tests serve canned responses.""" + return urllib.request.build_opener(_RefuseRedirects, *handlers) + + +_opener = _build_opener() + + +def _open_url(req: urllib.request.Request, timeout: float = 10): + """Open ``req`` through the redirect-refusing opener.""" + return _opener.open(req, timeout=timeout) + + +def is_civitai_url(url: str) -> bool: + """True when ``url`` is an HTTPS URL on an allow-listed CivitAI host.""" + try: + parts = urlsplit(url) + port = parts.port + except ValueError: + return False + hostname = (parts.hostname or "").lower() + return ( + parts.scheme == "https" + and hostname in CIVITAI_HOSTS + and parts.username is None + and parts.password is None + and port in (None, 443) + ) + + +def _read_capped(resp, cap: int) -> Optional[bytes]: + """Read a response body, refusing anything larger than ``cap`` bytes.""" + declared = resp.headers.get("Content-Length") + if declared is not None: + try: + if int(declared) > cap: + return None + except ValueError: + pass + + chunks = [] + total = 0 + while True: + chunk = resp.read(_DOWNLOAD_CHUNK_BYTES) + if not chunk: + break + total += len(chunk) + if total > cap: + return None + chunks.append(chunk) + return b"".join(chunks) + def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]: - """Download a single image, resize to thumbnail, save as JPEG.""" + """Download a single CivitAI image, resize to thumbnail, save as JPEG. + + Refuses any URL that is not HTTPS on an allow-listed CivitAI host so the + API key is never sent elsewhere, and refuses bodies larger than + MAX_CIVITAI_DOWNLOAD_BYTES. + """ + if not is_civitai_url(url): + logger.warning(f"Refusing non-CivitAI download URL: {url}") + return None try: headers = {"User-Agent": "ComfyUI-PromptManager/1.0"} if api_key: headers["Authorization"] = f"Bearer {api_key}" req = urllib.request.Request(url, headers=headers) - with urllib.request.urlopen(req, timeout=10) as resp: - raw = resp.read() + with _open_url(req, timeout=10) as resp: + raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES) + if raw is None: + logger.warning(f"Refusing oversized download: {url}") + return None # Resize to thumbnail to save disk space from io import BytesIO @@ -361,7 +538,7 @@ def download_civitai_images( if not isinstance(img, dict): continue url = img.get("url", "") - if not isinstance(url, str) or not url.startswith("http"): + if not isinstance(url, str) or not is_civitai_url(url): continue url_hash = hashlib.md5(url.encode()).hexdigest()[:12] diff --git a/pyproject.toml b/pyproject.toml index 188eef9..6cd4077 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,9 +1,14 @@ [project] name = "promptmanager" description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite." -version = "3.2.4" +version = "3.3.0" license = {file = "LICENSE"} -dependencies = ["# Core dependencies for PromptManager", "# Note: Most dependencies are already included with ComfyUI", "# Already included with Python standard library:", "# - sqlite3", "# - hashlib", "# - json", "# - datetime", "# - os", "# - typing", "# - threading", "# - uuid", "# Required for gallery functionality:", "watchdog>=2.1.0 # For file system monitoring", "Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)", "# Optional dependencies for enhanced search functionality:", "# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)", "# sqlalchemy>=1.4.0 # For advanced ORM features (optional)", "# Development dependencies (optional):", "# pytest>=6.0.0 # For running tests", "# black>=22.0.0 # For code formatting", "# flake8>=4.0.0 # For linting", "# mypy>=0.910 # For type checking"] +requires-python = ">=3.10" +dependencies = ["watchdog>=2.1.0", "Pillow>=8.0.0"] + +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" [project.urls] Repository = "https://github.com/ComfyAssets/ComfyUI_PromptManager" @@ -18,10 +23,5 @@ Icon = "https://avatars.githubusercontent.com/u/213204677?s=200" line-length = 88 target-version = ["py310"] -[tool.flake8] -max-line-length = 88 -extend-ignore = ["E203", "W503"] -exclude = ["reference/", ".git", "__pycache__", "web/", "venv/"] - [tool.bandit] exclude_dirs = ["tests", "reference"] diff --git a/requirements-dev.txt b/requirements-dev.txt index f2539a6..7338347 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -3,6 +3,10 @@ watchdog>=2.1.0 Pillow>=8.0.0 aiohttp>=3.8.0 +# Tests and coverage +coverage>=7.0 +numpy>=1.24 # WD14 engine tests + # Code quality black>=24.0.0 flake8>=7.0.0 diff --git a/requirements.txt b/requirements.txt index 815e196..f5f32b4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,26 +1,4 @@ -# Core dependencies for PromptManager -# Note: Most dependencies are already included with ComfyUI - -# Already included with Python standard library: -# - sqlite3 -# - hashlib -# - json -# - datetime -# - os -# - typing -# - threading -# - uuid - -# Required for gallery functionality: -watchdog>=2.1.0 # For file system monitoring -Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI) - -# Optional dependencies for enhanced search functionality: -# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional) -# sqlalchemy>=1.4.0 # For advanced ORM features (optional) - -# Development dependencies (optional): -# pytest>=6.0.0 # For running tests -# black>=22.0.0 # For code formatting -# flake8>=4.0.0 # For linting -# mypy>=0.910 # For type checking \ No newline at end of file +# Runtime dependencies for PromptManager. +# aiohttp, numpy and torch come from ComfyUI itself and are not listed here. +watchdog>=2.1.0 # Output directory monitoring +Pillow>=8.0.0 # Image metadata extraction and thumbnails diff --git a/restart_gallery.py b/restart_gallery.py deleted file mode 100644 index 7403594..0000000 --- a/restart_gallery.py +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env python3 -""" -ComfyUI_PromptManager Gallery System Restart Script. - -This script reinitializes and restarts the automatic image gallery monitoring system -for ComfyUI_PromptManager. Use this script after code modifications or when the -gallery system needs to be reloaded. - -The script performs the following operations: -1. Initializes the database connection -2. Creates a PromptTracker instance for linking prompts to images -3. Sets up the ImageMonitor for real-time file system monitoring -4. Starts monitoring ComfyUI output directories -5. Reports system status and readiness - -Usage: - python restart_gallery.py - -This is particularly useful during development when you need to restart the -gallery monitoring system without restarting the entire ComfyUI application. -""" - -import sys -import os - -# Add current directory to Python path -current_dir = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, current_dir) - -try: - # Import and initialize the gallery system - from database.operations import PromptDatabase - from utils.prompt_tracker import get_prompt_tracker - from utils.image_monitor import get_image_monitor - - print("[RESTART] Restarting PromptManager Gallery System...") - - # Initialize components using singleton getters - db = PromptDatabase() - tracker = get_prompt_tracker(db) - monitor = get_image_monitor(db, tracker) - - print("[SUCCESS] Components initialized successfully") - - # Test image monitoring directories - status = monitor.get_status() - print(f"[INFO] Monitoring status: {status}") - - # Start monitoring - monitor.start_monitoring() - print("[SUCCESS] Image monitoring restarted") - - print("\n[READY] Gallery system restart completed!") - print("Generate some images now to test the automatic linking.") - -except ImportError as e: - print(f"[ERROR] Import error: {e}") - print("Make sure you're running this from the PromptManager directory") - -except Exception as e: - print(f"[ERROR] Error: {e}") - import traceback - - traceback.print_exc() diff --git a/tests/js/api-client.test.js b/tests/js/api-client.test.js new file mode 100644 index 0000000..b226404 --- /dev/null +++ b/tests/js/api-client.test.js @@ -0,0 +1,168 @@ +// Run with: node --test tests/js/ +const test = require("node:test"); +const assert = require("node:assert/strict"); +const { createApiClient, ApiError, latestOnly, STALE } = require("../../web/js/api-client.js"); + +/** A fetch stand-in whose responses the test resolves by hand, in any order. */ +function fakeFetch() { + const calls = []; + const impl = (url, init = {}) => + new Promise((resolve, reject) => { + const call = { url, init, resolve, reject, aborted: false }; + if (init.signal) { + init.signal.addEventListener("abort", () => { + call.aborted = true; + reject(new DOMException("The operation was aborted.", "AbortError")); + }); + } + calls.push(call); + }); + impl.calls = calls; + return impl; +} + +function jsonResponse(status, body, { invalid = false } = {}) { + return { + ok: status >= 200 && status < 300, + status, + statusText: status === 200 ? "OK" : `Status ${status}`, + json: async () => { + if (invalid) throw new SyntaxError("Unexpected token"); + return body; + }, + }; +} + +test("get returns the parsed body of a successful response", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const pending = api.get("/prompt_manager/recent?page=1"); + fetch.calls[0].resolve(jsonResponse(200, { success: true, results: [1, 2] })); + assert.deepEqual(await pending, { success: true, results: [1, 2] }); + assert.equal(fetch.calls[0].init.method, "GET"); +}); + +test("get forwards an AbortSignal to fetch", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const controller = new AbortController(); + const pending = api.get("/x", { signal: controller.signal }); + assert.equal(fetch.calls[0].init.signal, controller.signal); + controller.abort(); + await assert.rejects(pending, (err) => err.name === "AbortError"); +}); + +test("non-2xx responses throw ApiError carrying the status and server message", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const pending = api.get("/x"); + fetch.calls[0].resolve(jsonResponse(404, { success: false, error: "Prompt not found" })); + await assert.rejects(pending, (err) => { + assert.ok(err instanceof ApiError); + assert.ok(err instanceof Error); + assert.equal(err.status, 404); + assert.equal(err.message, "Prompt not found"); + return true; + }); +}); + +test("non-2xx responses without a JSON body fall back to the status text", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const pending = api.get("/x"); + fetch.calls[0].resolve(jsonResponse(500, null, { invalid: true })); + await assert.rejects(pending, (err) => err.status === 500 && err.message === "Status 500"); +}); + +test("a 2xx body with success:false throws ApiError", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const pending = api.post("/x", {}); + fetch.calls[0].resolve(jsonResponse(200, { success: false, error: "Nope" })); + await assert.rejects(pending, (err) => err instanceof ApiError && err.status === 200 && err.message === "Nope"); +}); + +test("a 2xx body that is not JSON throws ApiError", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const pending = api.get("/x"); + fetch.calls[0].resolve(jsonResponse(200, null, { invalid: true })); + await assert.rejects(pending, (err) => err instanceof ApiError && err.status === 200); +}); + +test("post, put and del send JSON bodies with the right method", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const p1 = api.post("/a", { text: "x" }); + const p2 = api.put("/b", { rating: 3 }); + const p3 = api.del("/c"); + assert.equal(fetch.calls[0].init.method, "POST"); + assert.equal(fetch.calls[0].init.body, JSON.stringify({ text: "x" })); + assert.equal(fetch.calls[0].init.headers["Content-Type"], "application/json"); + assert.equal(fetch.calls[1].init.method, "PUT"); + assert.equal(fetch.calls[1].init.body, JSON.stringify({ rating: 3 })); + assert.equal(fetch.calls[2].init.method, "DELETE"); + assert.equal(fetch.calls[2].init.body, undefined); + for (const c of fetch.calls) c.resolve(jsonResponse(200, { success: true })); + await Promise.all([p1, p2, p3]); +}); + +test("network failures propagate as ordinary errors", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const pending = api.get("/x"); + fetch.calls[0].reject(new TypeError("Failed to fetch")); + await assert.rejects(pending, (err) => err instanceof TypeError && !(err instanceof ApiError)); +}); + +test("createApiClient without an argument uses the global fetch", async () => { + const original = globalThis.fetch; + globalThis.fetch = async () => jsonResponse(200, { success: true, via: "global" }); + try { + assert.deepEqual(await createApiClient().get("/x"), { success: true, via: "global" }); + } finally { + globalThis.fetch = original; + } +}); + +test("latestOnly: a newer call aborts the older one and the older result is STALE", async () => { + const fetch = fakeFetch(); + const api = createApiClient(fetch); + const latest = latestOnly(); + + const first = latest((signal) => api.get("/search?q=a", { signal })); + const second = latest((signal) => api.get("/search?q=ab", { signal })); + + assert.equal(fetch.calls[0].aborted, true, "first request aborted"); + assert.equal(fetch.calls[1].aborted, false); + + fetch.calls[1].resolve(jsonResponse(200, { success: true, q: "ab" })); + assert.deepEqual(await second, { success: true, q: "ab" }); + assert.equal(await first, STALE); +}); + +test("latestOnly: an older response that arrives after a newer request is STALE even without abort", async () => { + const latest = latestOnly(); + let resolveFirst; + const first = latest(() => new Promise((r) => { resolveFirst = r; })); + const second = latest(async () => "new"); + assert.equal(await second, "new"); + resolveFirst("old"); + assert.equal(await first, STALE); +}); + +test("latestOnly: sequential calls each return their own result", async () => { + const latest = latestOnly(); + assert.equal(await latest(async () => 1), 1); + assert.equal(await latest(async () => 2), 2); +}); + +test("latestOnly: errors from the current call propagate, errors from superseded calls do not", async () => { + const latest = latestOnly(); + let rejectFirst; + const first = latest(() => new Promise((_, rej) => { rejectFirst = rej; })); + const second = latest(async () => { throw new Error("boom"); }); + await assert.rejects(second, /boom/); + rejectFirst(new Error("old failure")); + assert.equal(await first, STALE); +}); diff --git a/tests/js/browser-globals.test.js b/tests/js/browser-globals.test.js new file mode 100644 index 0000000..135cde6 --- /dev/null +++ b/tests/js/browser-globals.test.js @@ -0,0 +1,74 @@ +// Run with: node --test tests/js/ +// Each page module is a UMD-style IIFE: CommonJS export under Node, a window +// global in the browser. This loads every module the way a browser would and +// checks the global name the pages rely on. +const test = require("node:test"); +const assert = require("node:assert/strict"); +const fs = require("node:fs"); +const path = require("node:path"); +const vm = require("node:vm"); + +const MODULES = { + "html-escape.js": ["HtmlEscape", "escapeHtml", "escapeAttr"], + "data-actions.js": ["DataActions"], + "api-client.js": ["ApiClient"], + "list-state.js": ["ListState"], + "sse-stream.js": ["SseStream"], + "png-metadata.js": ["PngMetadata"], + "view-router.js": ["ViewRouter"], + "prompt-list-sort.js": ["PromptListSort"], + "comfy-metadata.js": ["ComfyMetadata"], +}; + +function loadInFakeBrowser(file) { + const source = fs.readFileSync(path.join(__dirname, "../../web/js", file), "utf8"); + const window = {}; + window.window = window; + const context = vm.createContext({ + window, + globalThis: window, + console, + TextDecoder, + TextEncoder, + AbortController, + URLSearchParams, + Symbol, + Object, + Array, + Number, + String, + Math, + Set, + Map, + JSON, + Error, + Promise, + DataView, + Uint8Array, + ArrayBuffer, + }); + vm.runInContext(source, context, { filename: file }); + return window; +} + +for (const [file, globals] of Object.entries(MODULES)) { + test(`${file} exposes ${globals.join(", ")} on window when module is undefined`, () => { + const window = loadInFakeBrowser(file); + for (const name of globals) { + assert.ok(window[name], `${name} missing`); + assert.ok(["object", "function"].includes(typeof window[name]), `${name} is a ${typeof window[name]}`); + } + }); +} + +test("the browser escapeHtml global is the same function as HtmlEscape.escapeHtml", () => { + const window = loadInFakeBrowser("html-escape.js"); + assert.equal(window.escapeHtml, window.HtmlEscape.escapeHtml); + assert.equal(window.escapeAttr, window.escapeHtml); + assert.equal(window.escapeHtml(``), "<a href="x">"); +}); + +test("the browser ViewRouter global resolves the tags deep link", () => { + const window = loadInFakeBrowser("view-router.js"); + assert.equal(window.ViewRouter.resolveView("#/tags/a"), "tags"); +}); diff --git a/tests/js/data-actions.test.js b/tests/js/data-actions.test.js new file mode 100644 index 0000000..c449801 --- /dev/null +++ b/tests/js/data-actions.test.js @@ -0,0 +1,49 @@ +// Run with: node --test tests/js/ +const test = require("node:test"); +const assert = require("node:assert/strict"); +const { readAction, dispatch } = require("../../web/js/data-actions.js"); + +test("readAction parses promptId and index as integers", () => { + const parsed = readAction({ action: "open-film", promptId: "42", index: "3" }); + assert.equal(parsed.action, "open-film"); + assert.equal(parsed.promptId, 42); + assert.equal(parsed.index, 3); +}); + +test("readAction turns missing or non-numeric ids into null", () => { + assert.equal(readAction({ action: "edit" }).promptId, null); + assert.equal(readAction({ action: "edit", promptId: "abc" }).promptId, null); + assert.equal(readAction({ action: "edit", index: "" }).index, null); +}); + +test("readAction passes string fields through untouched", () => { + const parsed = readAction({ action: "remove-tag", tag: `a "quoted" tag`, url: "/x?y=1&z=2", copyText: "p" }); + assert.equal(parsed.tag, `a "quoted" tag`); + assert.equal(parsed.url, "/x?y=1&z=2"); + assert.equal(parsed.copyText, "p"); +}); + +test("readAction tolerates a missing dataset", () => { + assert.equal(readAction(undefined).action, ""); + assert.equal(readAction(null).promptId, null); +}); + +test("dispatch calls the matching handler with the parsed action and extras", () => { + const calls = []; + const handled = dispatch( + { action: "edit", promptId: "7" }, + { edit: (a, el) => calls.push([a.promptId, el]) }, + "element", + ); + assert.equal(handled, true); + assert.deepEqual(calls, [[7, "element"]]); +}); + +test("dispatch returns false and calls nothing for an unknown or empty action", () => { + let called = 0; + const handlers = { edit: () => called++ }; + assert.equal(dispatch({ action: "other" }, handlers), false); + assert.equal(dispatch({}, handlers), false); + assert.equal(dispatch({ action: "constructor" }, handlers), false, "no prototype lookups"); + assert.equal(called, 0); +}); diff --git a/tests/js/gallery-settings.test.js b/tests/js/gallery-settings.test.js new file mode 100644 index 0000000..c3681a2 --- /dev/null +++ b/tests/js/gallery-settings.test.js @@ -0,0 +1,63 @@ +// Run with: node --test tests/js/ +const test = require("node:test"); +const assert = require("node:assert/strict"); +const { + DEFAULT_SETTINGS, + normalizeGallerySettings, + gridColumnsStyle, + LIMIT_OPTIONS, +} = require("../../web/js/gallery-settings.js"); + +test("defaults are a complete, sane settings object", () => { + const s = normalizeGallerySettings(null); + assert.deepEqual(s, DEFAULT_SETTINGS); + assert.equal(s.defaultViewMode, "grid"); + assert.equal(s.defaultLimit, 100); + assert.equal(s.gridColumns, 8); + assert.equal(s.showImageInfo, true); + assert.equal(s.sidebarCollapsedByDefault, false); +}); + +test("stored values override defaults and unknown keys are kept", () => { + const s = normalizeGallerySettings({ defaultViewMode: "list", defaultLimit: 500, gridColumns: 10, showImageInfo: false, videoLoop: false, extra: 1 }); + assert.equal(s.defaultViewMode, "list"); + assert.equal(s.defaultLimit, 500); + assert.equal(s.gridColumns, 10); + assert.equal(s.showImageInfo, false); + assert.equal(s.videoLoop, false); + assert.equal(s.extra, 1); +}); + +test("invalid stored values fall back to defaults instead of breaking the page", () => { + const s = normalizeGallerySettings({ defaultViewMode: "mosaic", defaultLimit: 7, gridColumns: "lots", showImageInfo: "no" }); + assert.equal(s.defaultViewMode, "grid"); + assert.equal(s.defaultLimit, 100); + assert.equal(s.gridColumns, 8); + assert.equal(s.showImageInfo, true); +}); + +test("numeric strings from -
+
- + +
@@ -215,6 +227,11 @@ per page
+ +
Loading... @@ -401,6 +418,29 @@

Choose how the Web UI opens from ComfyUI nodes.

+ +
+
+ +

Load the next page of prompts as you scroll. The checkbox in the list header overrides this for the current session.

+
+ +
+ + + +
+

Performance

+ +
+
+ + 1 / 1 cores +
+ +

Threads used by the image rescan and thumbnail generation. More threads finish sooner; fewer leave CPU for ComfyUI while it generates.

+
@@ -442,6 +482,11 @@

Get your key from civitai.com/user/account

+ + +