From d5bcae301b0deab81f74e9d1944db2b6dd896054 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Fri, 23 Oct 2026 13:31:10 -0700 Subject: [PATCH] feat(system): run thumbnail generation as an Activity job GET /prompt_manager/images/thumbnails/plan counts what a run would do (targets, fresh, remaining past the cap, source labels) and writes nothing. POST /prompt_manager/jobs/thumbnails {mode, size, quality} starts the "thumbnails" job: admission guarded, cancellable at a file boundary, outcome groups generated / skipped / errors. size is the longest edge (128-2048), quality the encoder, and rebuild writes fresh previews again through the same atomic rename, so a failed file keeps its old preview and originals are only read. The legacy gallery stream now starts or attaches to the same job and streams its unchanged frames. A client that goes away no longer stops the work; the gallery's Cancel asks the server to cancel the job. --- TESTING.md | 25 ++ py/api/__init__.py | 3 + py/api/images.py | 297 +++++++-------- py/api/job_routes.py | 7 +- py/api/jobs.py | 29 ++ py/api/thumbnail_jobs.py | 448 +++++++++++++++++++++++ tests/js/gallery-thumbnails.test.js | 31 ++ tests/test_api_images.py | 32 +- tests/test_jobs.py | 22 ++ tests/test_thumbnail_jobs.py | 541 ++++++++++++++++++++++++++++ web/js/gallery.js | 22 +- 11 files changed, 1258 insertions(+), 199 deletions(-) create mode 100644 py/api/thumbnail_jobs.py create mode 100644 tests/test_thumbnail_jobs.py diff --git a/TESTING.md b/TESTING.md index 12adaece..d0c2272a 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1979,6 +1979,31 @@ Review options opens the legacy gallery in a new tab (D94) until its review dial A list that cannot be read shows "Sources could not be loaded. Your saved data has not changed." and Retry, never an empty list. +Thumbnail generation is an Activity job (S5, `py/api/thumbnail_jobs.py`). `GET +/prompt_manager/images/thumbnails/plan?mode=missing|rebuild&size=128..2048&quality=low|medium|high` +counts the work (`targets`, `fresh`, `remaining` past the 10,000-file cap, the source +labels) and writes nothing; a value out of range is a 400 `invalid-request` that starts +nothing. `POST /prompt_manager/jobs/thumbnails {mode, size, quality}` starts the +`thumbnails` job ("Thumbnail generation"): admission `guarded` (never beside a scan or +another record-changing job, one at a time, a legacy clear refused while it runs), +cancellable at a file boundary (a file not yet started is left alone; written previews +stay), outcome groups Generated / Skipped (already fresh, listed without a log line +each) / Errors. `size` is the longest edge and `quality` the encoder (low: JPEG 70, no +optimize; medium: 85, optimized, what the legacy routes always wrote; high: 95, +optimized). `rebuild` writes fresh previews again, each through the same temp-file +rename, so nothing is cleared first and a failed file keeps its old preview (a video +whose frame cannot be decoded keeps its preview instead of getting the placeholder); +originals are only read (`tests/test_thumbnail_jobs.py`). + +Deliberate change: the legacy gallery stream `POST /images/generate-thumbnails/progress` +now starts this job (missing previews, its quality preset as the size, medium encoding) +or attaches to the running one, and streams the job's unchanged `status`/`start`/ +`progress`/`file_error`/`complete` frames. A client that goes away no longer stops the +work (`test_thumbnail_job_survives_client_disconnect`); the gallery's Cancel asks the +server to cancel the running thumbnail job, and a job cancelled from Activity ends the +stream with the error frame "Thumbnail generation was cancelled". The blocking +`POST /images/generate-thumbnails` and the clear stay legacy holds. + Overview (`system-overview.mjs`) reads `GET /prompt_manager/system/overview` once per visit through `latestOnly()` on the section scope; a late answer after the section left paints nothing. Until it answers, the section shows "Loading the overview…"; a failure diff --git a/py/api/__init__.py b/py/api/__init__.py index 34ae5dac..b51a8db1 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -8,6 +8,7 @@ Split into domain-specific mixins for maintainability: - AutotagRoutesMixin: auto-tagging model management and tagging - JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py) - ImageTagJobsMixin: bulk Add tags / Auto-tag as image-tags jobs +- ThumbnailJobsMixin: thumbnail plan, thumbnail job and the gallery stream on it - SystemRoutesMixin: System health check, Overview and usage statistics """ @@ -37,6 +38,7 @@ from .admin import AdminRoutesMixin from .settings_routes import SettingsRoutesMixin from .job_routes import JobRoutesMixin from .image_tag_jobs import ImageTagJobsMixin +from .thumbnail_jobs import ThumbnailJobsMixin from .system import SystemRoutesMixin from .jobs import JobRegistry from .logging_routes import LoggingRoutesMixin @@ -272,6 +274,7 @@ class PromptManagerAPI( SettingsRoutesMixin, JobRoutesMixin, ImageTagJobsMixin, + ThumbnailJobsMixin, SystemRoutesMixin, LoggingRoutesMixin, AutotagRoutesMixin, diff --git a/py/api/images.py b/py/api/images.py index 3bf0a341..af8d60f4 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -20,7 +20,7 @@ from typing import Callable, NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError -from .jobs import THUMBNAILS_JOB_KIND +from .jobs import MAX_OUTCOME_ROWS, THUMBNAILS_JOB_KIND 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 @@ -235,6 +235,8 @@ MEDIA_EXTENSIONS = GALLERY_IMAGE_EXTENSIONS | VIDEO_EXTENSIONS # Caps on request-triggered filesystem work. MAX_SCAN_DEPTH = 12 MAX_THUMBNAIL_FILES = 10_000 +# Already-fresh files a thumbnail job lists under Skipped (its counter stays exact). +MAX_FRESH_LISTED = MAX_OUTCOME_ROWS def _worker_threads(): @@ -271,15 +273,31 @@ def _atomic_save(img, dst, fmt=None, **save_kwargs): tmp.unlink() -def _write_image_thumbnail(src, dst, thumbnail_size): - """Resize *src* into *dst* (blocking PIL work), atomically.""" +# Encoder settings per thumbnail quality (the System thumbnail dialog's +# "Medium · optimized", "High · original", "Low · fastest"). The legacy +# routes always encode as "medium", which is what they always wrote. +THUMBNAIL_ENCODINGS = { + "low": {"quality": 70, "optimize": False}, + "medium": {"quality": 85, "optimize": True}, + "high": {"quality": 95, "optimize": True}, +} +DEFAULT_THUMBNAIL_ENCODING = THUMBNAIL_ENCODINGS["medium"] + + +def _write_image_thumbnail(src, dst, thumbnail_size, encoding=None): + """Resize *src* into *dst* (blocking PIL work), atomically. + + ``encoding`` is one of THUMBNAIL_ENCODINGS (medium when None); a PNG + preview keeps only its ``optimize`` flag. + """ + encoding = encoding or DEFAULT_THUMBNAIL_ENCODING 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} + save_kwargs = dict(encoding) if dst.suffix.lower() == ".png": - save_kwargs = {"optimize": True} + save_kwargs = {"optimize": encoding["optimize"]} _atomic_save(img, dst, **save_kwargs) @@ -932,6 +950,10 @@ class ImageRoutesMixin: async def thumbnail_status_route(request): return await self.get_thumbnail_status(request) + @routes.get("/prompt_manager/images/thumbnails/plan") + async def thumbnail_plan_route(request): + return await self.get_thumbnail_plan(request) + @routes.post("/prompt_manager/images/generate-thumbnails") async def generate_thumbnails_route(request): return await self.generate_thumbnails(request) @@ -1804,15 +1826,30 @@ class ImageRoutesMixin: self.logger.info(payload["message"]) return payload - def _plan_gallery_thumbnails(self, output_dirs: list[Path]) -> dict: - """Plan every gallery root with one shared per-run generation cap.""" - combined = {"targets": [], "skipped": 0, "remaining": 0} + def _plan_gallery_thumbnails( + self, output_dirs: list[Path], rebuild: bool = False + ) -> dict: + """Plan every gallery root with one shared per-run generation cap. + + See ``_plan_thumbnails``; ``fresh_files`` keeps at most + MAX_FRESH_LISTED files across the roots. + """ + combined = { + "targets": [], + "skipped": 0, + "remaining": 0, + "fresh": 0, + "fresh_files": [], + } for root in output_dirs: root = Path(root) - plan = self._plan_thumbnails(root, root / "thumbnails") + plan = self._plan_thumbnails(root, root / "thumbnails", rebuild=rebuild) available = max(0, MAX_THUMBNAIL_FILES - len(combined["targets"])) combined["targets"].extend(plan["targets"][:available]) combined["skipped"] += plan["skipped"] + combined["fresh"] += plan["fresh"] + room = max(0, MAX_FRESH_LISTED - len(combined["fresh_files"])) + combined["fresh_files"].extend(plan["fresh_files"][:room]) combined["remaining"] += plan["remaining"] + max( 0, len(plan["targets"]) - available ) @@ -1845,21 +1882,26 @@ class ImageRoutesMixin: {"success": False, "error": _safe_error(exc)}, status=500 ) - def _plan_thumbnails(self, output_path, thumbnails_dir): + def _plan_thumbnails(self, output_path, thumbnails_dir, rebuild=False): """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. + Returns ``{"targets": [(src, dst, is_video)], "skipped": n, + "remaining": n, "fresh": n, "fresh_files": [src]}``. Files whose + thumbnail is already newer than the source are ``fresh``. Without + ``rebuild`` they are counted as skipped and never queued, so a large + library that is mostly done only pays for the files that changed; + ``fresh_files`` lists the first MAX_FRESH_LISTED of them. With + ``rebuild`` every file is a target (nothing is skipped). The walk is + bounded like the gallery listing (MAX_SCAN_DEPTH, MAX_GALLERY_FILES, + no symlinks) and at most MAX_THUMBNAIL_FILES targets are queued per + run; ``remaining`` counts the targets beyond that cap. """ output_path = Path(output_path) targets = [] skipped = 0 remaining = 0 + fresh = 0 + fresh_files = [] for src in _iter_media_files(output_path, MAX_GALLERY_FILES): suffix = src.suffix.lower() is_video = suffix in VIDEO_EXTENSIONS @@ -1870,8 +1912,12 @@ class ImageRoutesMixin: self.logger.warning(f"Skipping thumbnail outside safe dir: {src.name}") continue if _thumbnail_is_fresh(src, dst): - skipped += 1 - continue + fresh += 1 + if not rebuild: + skipped += 1 + if len(fresh_files) < MAX_FRESH_LISTED: + fresh_files.append(src) + continue if len(targets) >= MAX_THUMBNAIL_FILES: remaining += 1 continue @@ -1881,42 +1927,70 @@ class ImageRoutesMixin: 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} + return { + "targets": targets, + "skipped": skipped, + "remaining": remaining, + "fresh": fresh, + "fresh_files": fresh_files, + } 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): + def _generate_batch(self, targets, thumbnail_size, **options): """Create thumbnails for ``targets`` on the configured worker threads. - Blocking; returns one result dict per target, in order. + Blocking; returns one result dict per target, in order. ``options`` + go to ``_generate_one`` (``encoding``, ``rebuild``, ``stop``). """ def one(target): src, dst, is_video = target - return self._generate_one(src, dst, thumbnail_size, is_video) + return self._generate_one(src, dst, thumbnail_size, is_video, **options) 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.""" + def _generate_one( + self, + src, + dst, + thumbnail_size, + is_video, + encoding=None, + rebuild=False, + stop=None, + ): + """Create one thumbnail (blocking). Never raises; returns a result dict. + + ``stop()`` true before the file starts gives ``action: "cancelled"`` + and touches nothing (a job's file boundary). ``rebuild`` writes a + fresh preview again and never replaces an existing video preview with + the placeholder. Every write is atomic, so a failure keeps the old + preview. + """ result = { "file": src.name, "dir": src.parent.name, "type": "video" if is_video else "image", } try: + if stop is not None and stop(): + return {**result, "action": "cancelled"} if not src.is_file(): return {**result, "action": "error", "error": "File no longer exists"} - if _thumbnail_is_fresh(src, dst): + if not rebuild and _thumbnail_is_fresh(src, dst): 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): + made = self._generate_video_thumbnail( + src, dst, thumbnail_size, encoding, keep_existing=rebuild + ) + if not made: return {**result, "action": "error", "error": "Video decode failed"} else: - _write_image_thumbnail(src, dst, thumbnail_size) + _write_image_thumbnail(src, dst, thumbnail_size, encoding) return {**result, "action": "generated"} except UnidentifiedImageError: return {**result, "action": "error", "error": "Not a valid image file"} @@ -1924,147 +1998,6 @@ class ImageRoutesMixin: 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 (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", - }, - ) - busy = self._jobs.blocker(THUMBNAILS_JOB_KIND, admission="guarded") - if busy is not None: - # The gallery reads this stream: refuse with its error event. - await response.prepare(request) - message = f"{busy.title} already running" - await self._emit_sse( - response, "error", {"error": message, "message": message} - ) - return response - async with self._jobs.hold( - THUMBNAILS_JOB_KIND, THUMBNAIL_GENERATION_TITLE, admission="guarded" - ) as job: - await response.prepare(request) - if not await self._thumbnail_progress_stream(response, thumbnail_size): - job.finish("failed") # an error event ended the stream - return response - - async def _thumbnail_progress_stream(self, response, thumbnail_size) -> bool: - """Stream one generation run; False when it ended with an error event.""" - try: - output_dirs = self._get_all_output_dirs() - if not output_dirs: - await self._emit_sse( - response, "error", {"error": "ComfyUI output directory not found"} - ) - return False - - await self._emit_sse( - response, - "status", - {"phase": "scanning", "message": "Scanning for images and videos..."}, - ) - plan = await self._run_in_executor( - self._plan_gallery_thumbnails, output_dirs - ) - await self._stream_thumbnails( - response, - plan["targets"], - thumbnail_size, - skipped=plan["skipped"], - remaining=plan["remaining"], - ) - self.invalidate_gallery_cache() - return True - except Exception: - self.logger.exception("Thumbnail generation failed") - await self._emit_sse( - response, - "error", - { - "error": "An internal error occurred", - "message": "Thumbnail generation failed. Check server logs.", - }, - ) - return False - - 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: @@ -2075,8 +2008,22 @@ class ImageRoutesMixin: 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.""" + def _generate_video_thumbnail( + self, + video_path, + thumbnail_path, + thumbnail_size, + encoding=None, + keep_existing=False, + ): + """Generate thumbnail from video file. Returns True if successful. + + ``encoding`` is one of THUMBNAIL_ENCODINGS (medium when None). With + ``keep_existing`` an existing preview is kept, and the run fails, + when no frame could be decoded, rather than replaced by the + placeholder. + """ + encoding = encoding or DEFAULT_THUMBNAIL_ENCODING try: # Try using OpenCV first (most reliable) try: @@ -2100,7 +2047,7 @@ class ImageRoutesMixin: img = Image.fromarray(frame_rgb) img.thumbnail(thumbnail_size, Image.Resampling.LANCZOS) - _atomic_save(img, thumbnail_path, "JPEG", quality=85, optimize=True) + _atomic_save(img, thumbnail_path, "JPEG", **encoding) self.logger.debug( f"Generated video thumbnail using OpenCV: {thumbnail_path}" ) @@ -2148,6 +2095,12 @@ class ImageRoutesMixin: except (NameError, OSError): pass + if keep_existing and thumbnail_path.exists(): + self.logger.warning( + f"No frame decoded for {video_path.name}; keeping its preview" + ) + return False + # Last resort: create a placeholder thumbnail try: from PIL import ImageDraw, ImageFont diff --git a/py/api/job_routes.py b/py/api/job_routes.py index 83caa851..47056a11 100644 --- a/py/api/job_routes.py +++ b/py/api/job_routes.py @@ -1,7 +1,8 @@ """Session job routes (Activity): list, inspect, follow and cancel jobs. The registry itself is ``py/api/jobs.py``. The scan job's runner is in -``py/api/admin.py``, the image tagging job's in ``py/api/image_tag_jobs.py``. +``py/api/admin.py``, the image tagging job's in ``py/api/image_tag_jobs.py`` +and the thumbnail job's in ``py/api/thumbnail_jobs.py``. Every response leaves out absolute server paths: job rows already carry root-relative paths and public root labels. """ @@ -129,6 +130,10 @@ class JobRoutesMixin: async def start_image_tag_job_route(request): return await self.start_image_tag_job(request) + @routes.post("/prompt_manager/jobs/thumbnails") + async def start_thumbnail_job_route(request): + return await self.start_thumbnail_job(request) + @routes.get("/prompt_manager/jobs/{job_id}") async def job_detail_route(request): return await self.job_detail(request) diff --git a/py/api/jobs.py b/py/api/jobs.py index 09ab15aa..64515e77 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -72,6 +72,9 @@ SCAN_GROUPS = { IMAGE_TAG_JOB_KIND = "image-tags" IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors") +# Thumbnail generation (#249 System S5): previews written, already fresh, failed. +THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors") + ADMISSIONS = ("read", "guarded", "records") # What each admission class is refused beside (besides its own kind). _REFUSED_BESIDE = { @@ -247,6 +250,32 @@ class Job: ) self.log_line(log_line) + def record_counted(self, group, count, rows=()): + """Count ``count`` outcomes in ``group`` at once, listing ``rows``. + + For outcomes that need no log line each (a thumbnail job's already + fresh previews): the counter is exact, ``rows`` (``record``'s row + shape) are listed up to MAX_OUTCOME_ROWS, and nothing is logged. + """ + if group not in self.counters: + raise ValueError(f"Unknown outcome group {group!r} for {self.kind}") + self.counters[group] += max(0, int(count)) + listed = self.outcomes[group] + for row in rows: + if len(listed) >= MAX_OUTCOME_ROWS: + break + index = row.get("root_index") + if isinstance(index, bool) or not isinstance(index, int) or index < 0: + index = None + listed.append( + { + "path": _relative_path(row.get("path")), + "root": str(row.get("root") or ""), + "root_index": index, + "detail": str(row.get("detail") or ""), + } + ) + def finish(self, state="finished"): """End the job in ``state`` and release every subscriber. diff --git a/py/api/thumbnail_jobs.py b/py/api/thumbnail_jobs.py new file mode 100644 index 00000000..b75ab8b3 --- /dev/null +++ b/py/api/thumbnail_jobs.py @@ -0,0 +1,448 @@ +"""Thumbnail generation as an Activity job (#249 System S5, Task 12). + +- ``GET /prompt_manager/images/thumbnails/plan?mode=&size=&quality=`` counts + what a run would do over every gallery root and writes nothing: + ``{success, targets, fresh, remaining, roots, mode, size, quality}``. +- ``POST /prompt_manager/jobs/thumbnails {mode, size, quality}`` starts the + ``thumbnails`` job (202 ``{job}``). Its admission class is ``guarded``: it + never runs beside a record-changing job, and one thumbnail job runs at a + time. It stops at a file boundary when cancelled, keeping the previews it + wrote. Its outcome groups are ``generated``, ``skipped`` (already fresh) + and ``errors``. + +``mode`` is ``missing`` (previews that are absent or older than their file) +or ``rebuild`` (every preview again, fresh ones too). ``size`` is the +preview's longest edge, 128 to 2048 pixels. ``quality`` picks the encoder +settings (``THUMBNAIL_ENCODINGS``). Every preview is written next to the old +one and renamed over it, so nothing is cleared first and a failure keeps the +old preview; originals are only read. + +The legacy gallery stream ``POST /images/generate-thumbnails/progress`` +starts this job (its quality preset still sizes the previews) or attaches to +the running one, and streams the job's legacy ``status``/``start``/ +``progress``/``file_error``/``complete`` frames. A client that goes away no +longer stops the work: the job outlives the request, as a scan does. +""" + +import asyncio +import functools +import json +import time +from dataclasses import dataclass + +from aiohttp import web + +from .admin import _plural, _scan_location, _scan_roots, _scan_status_line +from .images import ( + THUMBNAIL_ENCODINGS, + _BodyTooLarge, + _new_thumbnail_stats, + _read_json_body, + _record_thumbnail_result, + _thumbnail_batch_size, + _thumbnail_complete_payload, + _thumbnail_progress_payload, + _worker_threads, +) +from .job_routes import busy_response +from .jobs import THUMBNAIL_GROUP_NAMES, THUMBNAILS_JOB_KIND, JobBusy + +THUMBNAIL_JOB_TITLE = "Thumbnail generation" +THUMBNAIL_MODES = ("missing", "rebuild") +THUMBNAIL_QUALITIES = tuple(THUMBNAIL_ENCODINGS) +MIN_THUMBNAIL_EDGE = 128 +MAX_THUMBNAIL_EDGE = 2048 +# The concept's "Preview size limit" default. +DEFAULT_THUMBNAIL_EDGE = 512 +DEFAULT_THUMBNAIL_OPTIONS = { + "mode": "missing", + "size": DEFAULT_THUMBNAIL_EDGE, + "quality": "medium", +} +_OPTIONS_ERROR = ( + 'mode must be "missing" or "rebuild", quality "low", "medium" or "high", ' + f"and size a whole number of pixels from {MIN_THUMBNAIL_EDGE} " + f"to {MAX_THUMBNAIL_EDGE}" +) +CANCELLED_ERROR = "Thumbnail generation was cancelled" +_MODE_LOG = {"missing": "missing previews", "rebuild": "every preview"} +# Outcome rows: what happened to each file's preview. +_DETAILS = { + "generated": "Preview written", + "skipped": "Preview already up to date", +} +_LOG_PREFIXES = {"generated": "+", "errors": "!"} + + +@dataclass(frozen=True) +class ThumbnailOptions: + """One run's choices: the mode, the preview's longest edge and quality.""" + + mode: str + edge: int + quality: str + + @property + def size(self): + return (self.edge, self.edge) + + @property + def encoding(self): + return THUMBNAIL_ENCODINGS[self.quality] + + @property + def rebuild(self): + return self.mode == "rebuild" + + def scope(self): + return {"mode": self.mode, "size": self.edge, "quality": self.quality} + + +def _edge(value, *, from_query): + """The preview edge as an int, or None when it is not one in range.""" + if from_query and isinstance(value, str): + if not value.isascii() or not value.isdigit() or len(value) > 5: + return None + value = int(value) + if isinstance(value, bool) or not isinstance(value, int): + return None + return value if MIN_THUMBNAIL_EDGE <= value <= MAX_THUMBNAIL_EDGE else None + + +def thumbnail_options(raw, *, from_query=False): + """``(ThumbnailOptions, None)`` from a body or query, or ``(None, message)``. + + Absent fields take ``DEFAULT_THUMBNAIL_OPTIONS``. Query values are + strings, so a query ``size`` may be decimal digits; a JSON ``size`` must + be an int. + """ + if not isinstance(raw, dict) and not hasattr(raw, "get"): + return None, _OPTIONS_ERROR + mode = raw.get("mode", DEFAULT_THUMBNAIL_OPTIONS["mode"]) + quality = raw.get("quality", DEFAULT_THUMBNAIL_OPTIONS["quality"]) + edge = _edge(raw.get("size", DEFAULT_THUMBNAIL_EDGE), from_query=from_query) + if mode not in THUMBNAIL_MODES or quality not in THUMBNAIL_QUALITIES: + return None, _OPTIONS_ERROR + if edge is None: + return None, _OPTIONS_ERROR + return ThumbnailOptions(mode, edge, quality), None + + +def _invalid(message): + return web.json_response( + {"success": False, "code": "invalid-request", "error": message}, status=400 + ) + + +def _frame(event, data): + """A legacy stream frame as the job publishes it.""" + return {"event": event, "data": data} + + +class ThumbnailJobsMixin: + """The thumbnail plan, the thumbnail job and the legacy stream on it.""" + + async def get_thumbnail_plan(self, request): + """What a thumbnail run would do (read-only; capped like a run).""" + options, error = thumbnail_options(request.query, from_query=True) + if error: + return _invalid(error) + try: + output_dirs = await self._run_in_executor(self._get_all_output_dirs) + plan = await self._run_in_executor( + self._plan_gallery_thumbnails, output_dirs, rebuild=options.rebuild + ) + except Exception: + self.logger.exception("Thumbnail plan failed") + return web.json_response( + { + "success": False, + "code": None, + "error": "The thumbnail plan could not be read", + }, + status=500, + ) + return web.json_response( + { + "success": True, + "targets": len(plan["targets"]), + "fresh": plan["fresh"], + "remaining": plan["remaining"], + "roots": [self._public_path(root) for root in output_dirs], + **options.scope(), + }, + headers={"Cache-Control": "no-store"}, + ) + + async def start_thumbnail_job(self, request): + """Start the thumbnail job (202), or 409 naming the job in the way.""" + body = {} + if request.can_read_body: + try: + body = await _read_json_body(request) + except _BodyTooLarge: + return web.json_response( + {"success": False, "code": "too-large", "error": "Body too large"}, + status=413, + ) + except (ValueError, UnicodeDecodeError): + return _invalid(_OPTIONS_ERROR) + options, error = thumbnail_options(body) + if error: + return _invalid(error) + try: + job = self._start_thumbnail_job(options) + except JobBusy as busy: + return busy_response(busy) + return web.json_response({"success": True, "job": job.summary()}, status=202) + + def _start_thumbnail_job(self, options): + """Register the thumbnail job synchronously; raises JobBusy.""" + return self._jobs.start( + THUMBNAILS_JOB_KIND, + THUMBNAIL_JOB_TITLE, + admission="guarded", + cancellable=True, + scope=options.scope(), + runner=functools.partial(self._run_thumbnail_job, options=options), + groups=THUMBNAIL_GROUP_NAMES, + ) + + # ── the job ─────────────────────────────────────────────────────── + + async def _run_thumbnail_job(self, job, options): + try: + state = await self._thumbnail_job_body(job, options) + except Exception: + self.logger.exception("Thumbnail generation failed") + job.update(status_line="Thumbnail generation stopped on an internal error.") + job.log_line("Thumbnails failed · internal error") + job.publish( + _frame( + "error", + { + "error": "An internal error occurred", + "message": "Thumbnail generation failed. Check server logs.", + }, + ) + ) + state = "failed" + job.finish(state) + + async def _thumbnail_job_body(self, job, options): + """Plan, then generate in worker batches; returns the final state.""" + output_dirs = await self._run_in_executor(self._get_all_output_dirs) + if not output_dirs: + message = "ComfyUI output directory not found" + job.update(status_line=message) + job.log_line(f"Thumbnails failed · {message}") + job.publish(_frame("error", {"error": message})) + return "failed" + folders = _plural(len(output_dirs), "output folder") + job.scope["output_folders"] = len(output_dirs) + job.update(status_line=f"Looking for media files in {folders}") + job.log_line( + f"Thumbnails started · {folders} · {_MODE_LOG[options.mode]} · " + f"{options.edge} px · {options.quality}" + ) + job.publish( + _frame( + "status", + {"phase": "scanning", "message": "Scanning for images and videos..."}, + ) + ) + plan = await self._run_in_executor( + self._plan_gallery_thumbnails, output_dirs, rebuild=options.rebuild + ) + roots = _scan_roots(output_dirs) + self._record_fresh(job, plan, roots) + stats, start = await self._generate_thumbnail_targets(job, options, plan, roots) + self.invalidate_gallery_cache() + generated = _plural(stats["generated"], "preview") + errors = _plural(len(stats["errors"]), "error") + if job.cancelled: + job.log_line(f"Thumbnails cancelled · {generated} written") + job.publish( + _frame("error", {"error": CANCELLED_ERROR, "message": CANCELLED_ERROR}) + ) + return "cancelled" + payload = _thumbnail_complete_payload( + len(plan["targets"]) + plan["skipped"] + plan["remaining"], + stats, + start, + remaining=plan["remaining"], + processed=len(plan["targets"]), + ) + job.update(progress=100) + if plan["remaining"]: + job.log_line( + f"{plan['remaining']:,} more previews are left for the next run" + ) + job.log_line(f"Thumbnails finished · {generated} written · {errors}") + self.logger.info(payload["message"]) + job.publish(_frame("complete", payload)) + return "finished" + + def _record_fresh(self, job, plan, roots): + """List the already fresh previews under Skipped (no log line each).""" + rows = [] + for src in plan["fresh_files"]: + path, label, index = _scan_location(src, roots) + rows.append( + { + "path": path, + "root": label, + "root_index": index, + "detail": _DETAILS["skipped"], + } + ) + job.record_counted("skipped", plan["skipped"], rows) + if plan["skipped"]: + job.log_line(f"= {plan['skipped']:,} previews already up to date") + + async def _generate_thumbnail_targets(self, job, options, plan, roots): + """Generate ``plan``'s targets in worker batches, publishing each file.""" + targets = plan["targets"] + 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 plan["skipped"]: + message += f" ({plan['skipped']} already have thumbnails)" + job.update(total=total, status_line=_scan_status_line(0, total, 0, 0)) + job.publish( + _frame( + "start", + { + "total_images": total, + "skipped": plan["skipped"], + "remaining": plan["remaining"], + "phase": "processing", + "image_count": total - video_count, + "video_count": video_count, + "message": message, + }, + ) + ) + stats = _new_thumbnail_stats() + stats["skipped"] = plan["skipped"] + start = time.monotonic() + batch_size = _thumbnail_batch_size(_worker_threads()) + generate = { + "encoding": options.encoding, + "rebuild": options.rebuild, + "stop": lambda: job.cancelled, + } + index = 0 + batches = 0 + for batch_start in range(0, total, batch_size): + if job.cancelled: + break + batch = targets[batch_start : batch_start + batch_size] + results = await self._run_in_executor( + self._generate_batch, batch, options.size, **generate + ) + for (src, _dst, _video), result in zip(batch, results): + if result["action"] == "cancelled": + continue + index += 1 + _record_thumbnail_result(stats, result) + self._record_thumbnail_outcome(job, src, result, roots) + job.update(processed=index, progress=index * 100 // total) + if result["action"] == "error": + job.publish( + _frame( + "file_error", + {"file": result["file"], "error": result["error"]}, + ) + ) + job.publish( + _frame( + "progress", + _thumbnail_progress_payload(index, total, stats, result, start), + ) + ) + batches += 1 + job.update( + status_line=_scan_status_line( + index, total, batches, time.monotonic() - start + ) + ) + await asyncio.sleep(0) # let other requests run between batches + return stats, start + + @staticmethod + def _record_thumbnail_outcome(job, src, result, roots): + path, label, index = _scan_location(src, roots) + where = f"{label}/{path}" if label else path + if result["action"] == "error": + group, detail = "errors", result.get("error") or "failed" + elif result["action"] == "skipped": + group, detail = "skipped", _DETAILS["skipped"] + else: + group, detail = "generated", _DETAILS["generated"] + row = {"path": path, "root": label, "root_index": index, "detail": detail} + if group == "skipped": + job.record_counted(group, 1, [row]) + return + job.record(group, row, f"{_LOG_PREFIXES[group]} {where} · {detail}") + + # ── the legacy gallery stream ───────────────────────────────────── + + async def generate_thumbnails_with_progress(self, request): + """The gallery's thumbnail stream (POST, server-sent events). + + Starts the thumbnail job with the gallery's quality preset as the + preview size (missing previews, medium encoding, as before), or + attaches to the running thumbnail job, and streams its frames. A + legacy hold (the blocking generate or clear) or a record-changing + job refuses it with the stream's error frame. The job goes on when + the client goes away. + """ + 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", + }, + ) + running = self._jobs.active(THUMBNAILS_JOB_KIND) + job = running if running is not None and running.cancellable else None + busy = None + if job is None: + busy = self._jobs.blocker(THUMBNAILS_JOB_KIND, admission="guarded") + await response.prepare(request) + if busy is None and job is None: + # The gallery's preset sizes the previews; medium encodes them as + # the gallery always did. + options = ThumbnailOptions("missing", thumbnail_size[0], "medium") + try: + job = self._start_thumbnail_job(options) + except JobBusy as refused: + busy = refused.running + if busy is not None: + message = f"{busy.title} already running" + await self._emit_sse( + response, "error", {"error": message, "message": message} + ) + return response + # No await between the start and the first read: the stream subscribes + # before the job's runner publishes its first frame. + events = job.events() + try: + async for payload in events: + frame = ( + f"event: {payload['event']}\ndata: {json.dumps(payload['data'])}" + ) + await response.write(f"{frame}\n\n".encode("utf-8")) + except Exception as exc: # the client went away; the job goes on + self.logger.info( + f"Thumbnail progress client disconnected ({exc}); " + "the thumbnails continue in the background" + ) + finally: + await events.aclose() + return response diff --git a/tests/js/gallery-thumbnails.test.js b/tests/js/gallery-thumbnails.test.js index 5b085f21..472c2927 100644 --- a/tests/js/gallery-thumbnails.test.js +++ b/tests/js/gallery-thumbnails.test.js @@ -110,3 +110,34 @@ test('saving gallery settings retains the remembered thumbnail dismissal', () => f.manager.saveSettings(); assert.equal(JSON.parse(f.storage.get('gallerySettings')).thumbnailPromptDismissed, 'remembered-state'); }); + +// #249 System S5: the stream's work is now an Activity job that outlives the page's connection, so the gallery's +// Cancel asks the server to stop the running thumbnail job instead of only closing the stream. +test('Cancel stops the running thumbnail job on the server', async () => { + const f = fixture({}); + const posts = []; + const running = { id: 'b'.repeat(32), kind: 'thumbnails', state: 'running', cancellable: true }; + f.manager.api = { + get: async url => { f.requests.push(url); return { success: true, jobs: [ + { id: 'a'.repeat(32), kind: 'scan', state: 'running', cancellable: true }, + { id: 'c'.repeat(32), kind: 'thumbnails', state: 'finished', cancellable: true }, running] }; }, + post: async (url, body) => { posts.push([url, body]); return { success: true, cancellation_requested: true }; }, + }; + assert.equal(await f.manager.cancelThumbnailJob(), true); + assert.deepEqual(f.requests, ['/prompt_manager/jobs']); + assert.equal(JSON.stringify(posts), JSON.stringify([[`/prompt_manager/jobs/${running.id}/cancel`, {}]])); +}); + +test('Cancel posts nothing when no cancellable thumbnail job runs, and survives a failed read', async () => { + const f = fixture({}); + const posts = []; + f.manager.api = { + get: async () => ({ success: true, jobs: [{ id: 'd'.repeat(32), kind: 'thumbnails', state: 'running', + cancellable: false }] }), + post: async url => { posts.push(url); return {}; }, + }; + assert.equal(await f.manager.cancelThumbnailJob(), false); + f.manager.api.get = async () => { throw new Error('offline'); }; + assert.equal(await f.manager.cancelThumbnailJob(), false); + assert.deepEqual(posts, []); +}); diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 8c0361cc..afe37cc7 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -413,11 +413,11 @@ class TestThumbnailProgressStream(ThumbnailTestCase): lock = threading.Lock() original = images_module._write_image_thumbnail - def recording(src, dst, size): + def recording(src, dst, size, *args): with lock: seen.add(threading.current_thread().name) time.sleep(0.05) - return original(src, dst, size) + return original(src, dst, size, *args) images_module._write_image_thumbnail = recording self.addCleanup(setattr, images_module, "_write_image_thumbnail", original) @@ -1939,7 +1939,11 @@ class TestGalleryRescan(ImageAPITestCase): class TestOutputScanHelpers(ImageAPITestCase): - """Stable ids for filesystem entries, no symlink reads, stop on disconnect.""" + """Stable ids for filesystem entries and no symlink reads. + + The thumbnail stream no longer stops when its client goes away: see + ``tests/test_thumbnail_jobs.py`` (the work is a job that outlives it). + """ def test_output_entry_id_is_a_stable_digest_of_the_path(self): import hashlib @@ -1964,28 +1968,6 @@ class TestOutputScanHelpers(ImageAPITestCase): names = sorted(p.name for p in images_module._iter_media_files(output, 100)) self.assertEqual(names, ["real.png"]) - async def test_thumbnail_stream_stops_when_the_client_disconnects(self): - generated = [] - - def fake_generate(src, dst, size, is_video): - generated.append(src) - return { - "action": "generated", - "file": src.name, - "dir": src.parent.name, - "type": "image", - } - - self.api._generate_one = fake_generate - - class GoneResponse: - async def write(self, data): - raise ConnectionResetError("client went away") - - targets = [(Path(f"/x/{i}.png"), Path(f"/x/t{i}.jpg"), False) for i in range(3)] - await self.api._stream_thumbnails(GoneResponse(), targets, (256, 256)) - self.assertEqual(generated, []) - class TestThumbnailStartupStatus(ThumbnailTestCase): """Startup counts and generation cover the same configured collection.""" diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 6848fa8a..59ba655e 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -421,6 +421,28 @@ class TestJobOutcomesAndLog(unittest.TestCase): SCAN_GROUP_NAMES, ("added", "known", "review", "no_metadata", "errors") ) + def test_record_counted_counts_exactly_lists_capped_and_logs_nothing(self): + job = Job("thumbnails", "Thumbnail generation", groups=("skipped",)) + rows = [ + {"path": os.path.abspath(f"out{os.sep}f{i}.png"), "root": "output"} + for i in range(MAX_OUTCOME_ROWS + 5) + ] + job.record_counted("skipped", 10_000, rows) + job.record_counted("skipped", 1, [{"path": "late.png", "root_index": True}]) + self.assertEqual(job.counters["skipped"], 10_001) + self.assertEqual(len(job.outcomes["skipped"]), MAX_OUTCOME_ROWS) + self.assertEqual(job.outcomes["skipped"][0]["path"], "f0.png") + self.assertIsNone(job.outcomes["skipped"][0]["root_index"]) + self.assertEqual(job.log, []) + with self.assertRaises(ValueError): + job.record_counted("bogus", 1) + + def test_thumbnail_job_groups(self): + self.assertEqual( + jobs_module.THUMBNAIL_GROUP_NAMES, ("generated", "skipped", "errors") + ) + self.assertEqual(jobs_module.THUMBNAILS_JOB_KIND, "thumbnails") + def test_image_tag_job_kind_and_groups(self): # #249 D1: bulk Add tags / Auto-tag, counter-tile order. self.assertEqual(getattr(jobs_module, "IMAGE_TAG_JOB_KIND", None), "image-tags") diff --git a/tests/test_thumbnail_jobs.py b/tests/test_thumbnail_jobs.py new file mode 100644 index 00000000..a602b6a3 --- /dev/null +++ b/tests/test_thumbnail_jobs.py @@ -0,0 +1,541 @@ +"""Thumbnail generation as an Activity job (#249 System S5, Task 12). + +``GET /prompt_manager/images/thumbnails/plan`` counts what a run would do and +writes nothing. ``POST /prompt_manager/jobs/thumbnails`` starts the +``thumbnails`` job (admission ``guarded``, cancellable at a file boundary). +The legacy gallery stream ``POST /images/generate-thumbnails/progress`` starts +or attaches to that job and streams its unchanged events from it, so a client +that goes away no longer stops the work. + +Ordering is driven by threading/asyncio events, never by timers. +""" + +import asyncio +import hashlib +import os +import sys +import threading +import unittest +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +try: + from tests.test_api_images import ThumbnailTestCase, make_png, parse_sse +except ImportError: # discovered with tests/ as the top-level directory + from test_api_images import ThumbnailTestCase, make_png, parse_sse + +from PIL import Image # noqa: E402 + +import py.api.images as images_module # noqa: E402 +from py.config import PromptManagerConfig # noqa: E402 + +PLAN_URL = "/prompt_manager/images/thumbnails/plan" +JOB_URL = "/prompt_manager/jobs/thumbnails" + + +def _tree_state(root): + """{relative path: (size, mtime_ns, sha256)} of every file under ``root``.""" + state = {} + for folder, _dirs, files in os.walk(root): + for name in files: + path = Path(folder) / name + data = path.read_bytes() + stat = path.stat() + state[path.relative_to(root).as_posix()] = ( + stat.st_size, + stat.st_mtime_ns, + hashlib.sha256(data).hexdigest(), + ) + return state + + +class ThumbnailJobCase(ThumbnailTestCase): + """Helpers: start, gate and wait for thumbnail jobs.""" + + async def _start(self, **body): + resp = await self.client.post(JOB_URL, json=body) + return resp.status, await resp.json() + + async def _finished(self, job_id): + job = self.api._jobs.get(job_id) + await asyncio.wait_for(job.task, 10) + return job + + def _gate_generation(self): + """Block every file in the worker until released; returns (entered, release).""" + entered = threading.Event() + release = threading.Event() + original = self.api._generate_one + + def gated(*args, **kwargs): + entered.set() + release.wait(10) + return original(*args, **kwargs) + + self.api._generate_one = gated + self.addCleanup(release.set) + return entered, release + + async def _wait(self, event): + loop = asyncio.get_running_loop() + self.assertTrue(await loop.run_in_executor(None, event.wait, 10)) + + def _run_scan(self): + gate = asyncio.Event() + + async def runner(job): + await gate.wait() + + job = self.api._jobs.start( + "scan", "Scan for metadata", admission="records", scope={}, runner=runner + ) + + async def release(): + gate.set() + await asyncio.wait_for(job.task, 5) + + self.addAsyncCleanup(release) + return job + + +class TestThumbnailPlan(ThumbnailJobCase): + async def test_plan_is_read_only_and_counts_cap(self): + for name in ("a", "b", "c"): + make_png(self.output_dir / f"{name}.png", size=(64, 64)) + fresh = make_png(self._thumb_path("a")) + source_time = (self.output_dir / "a.png").stat().st_mtime + os.utime(fresh, (source_time + 10, source_time + 10)) + before = _tree_state(self.comfy_dir) + + with patch.object(images_module, "MAX_THUMBNAIL_FILES", 1): + missing = await self.client.get( + PLAN_URL, params={"mode": "missing", "size": "512", "quality": "high"} + ) + rebuild = await self.client.get( + PLAN_URL, params={"mode": "rebuild", "size": "256", "quality": "low"} + ) + + self.assertEqual(missing.status, 200) + self.assertEqual( + await missing.json(), + { + "success": True, + "targets": 1, + "fresh": 1, + "remaining": 1, + "roots": ["output"], + "mode": "missing", + "size": 512, + "quality": "high", + }, + ) + data = await rebuild.json() + # Rebuild regenerates fresh previews too: all three are work, one fits the cap. + self.assertEqual((data["targets"], data["fresh"], data["remaining"]), (1, 1, 2)) + self.assertEqual( + (data["mode"], data["size"], data["quality"]), ("rebuild", 256, "low") + ) + self.assertEqual(_tree_state(self.comfy_dir), before) + self.assertEqual(self.api._jobs.list(), []) + + async def test_plan_defaults_to_missing_previews_at_the_concept_size(self): + make_png(self.output_dir / "a.png") + data = await (await self.client.get(PLAN_URL)).json() + self.assertEqual( + (data["mode"], data["size"], data["quality"], data["targets"]), + ("missing", 512, "medium", 1), + ) + + async def test_size_bounds_validated(self): + make_png(self.output_dir / "a.png") + for params in ( + {"size": "127"}, + {"size": "2049"}, + {"size": "512.5"}, + {"size": "big"}, + {"mode": "all"}, + {"quality": "ultra"}, + ): + with self.subTest(params=params): + resp = await self.client.get(PLAN_URL, params=params) + body = await resp.json() + self.assertEqual(resp.status, 400) + self.assertEqual( + (body["success"], body["code"]), (False, "invalid-request") + ) + for body in ( + {"size": 127}, + {"size": 2049}, + {"size": "512"}, + {"size": 512.0}, + {"size": True}, + {"mode": "rebuild-all"}, + {"quality": None}, + [], + ): + with self.subTest(body=body): + resp = await self.client.post(JOB_URL, json=body) + self.assertEqual(resp.status, 400) + self.assertEqual((await resp.json())["code"], "invalid-request") + for edge in (128, 2048): + with self.subTest(edge=edge): + resp = await self.client.get(PLAN_URL, params={"size": str(edge)}) + self.assertEqual(resp.status, 200) + self.assertEqual(self.api._jobs.list(), []) + self.assertFalse(self._thumb_path("a").exists()) + + +class TestThumbnailJob(ThumbnailJobCase): + async def test_job_generates_missing_previews_at_the_chosen_edge(self): + make_png(self.output_dir / "a.png", size=(800, 400)) + make_png(self.output_dir / "sub" / "b.png", size=(64, 64)) + (self.output_dir / "bad.png").write_bytes(b"\x89PNG not really") + + status, body = await self._start(mode="missing", size=200, quality="medium") + + self.assertEqual(status, 202) + summary = body["job"] + self.assertEqual( + (summary["kind"], summary["title"], summary["admission"]), + ("thumbnails", "Thumbnail generation", "guarded"), + ) + self.assertTrue(summary["cancellable"]) + self.assertEqual(set(summary["counters"]), {"generated", "skipped", "errors"}) + self.assertEqual( + summary["scope"], {"mode": "missing", "size": 200, "quality": "medium"} + ) + job = await self._finished(summary["id"]) + self.assertEqual(job.state, "finished") + self.assertEqual(job.counters, {"generated": 2, "skipped": 0, "errors": 1}) + self.assertEqual((job.processed, job.total, job.progress), (3, 3, 100)) + self.assertEqual(job.scope["output_folders"], 1) + with Image.open(self._thumb_path("a")) as thumb: + self.assertEqual(thumb.size, (200, 100)) + rows = {row["path"]: row for row in job.outcomes["generated"]} + self.assertEqual(set(rows), {"a.png", "sub/b.png"}) + self.assertEqual( + (rows["a.png"]["root"], rows["a.png"]["root_index"]), ("output", 0) + ) + self.assertEqual(job.outcomes["errors"][0]["path"], "bad.png") + self.assertTrue(any("Thumbnails finished" in line for line in job.log)) + # Rows and log lines never carry the server's absolute layout. + self.assertNotIn(str(self.comfy_dir), repr(job.detail())) + + async def test_fresh_previews_are_counted_as_skipped(self): + make_png(self.output_dir / "a.png") + make_png(self.output_dir / "b.png") + fresh = make_png(self._thumb_path("a")) + source_time = (self.output_dir / "a.png").stat().st_mtime + os.utime(fresh, (source_time + 10, source_time + 10)) + + _, body = await self._start() + job = await self._finished(body["job"]["id"]) + + self.assertEqual(job.counters, {"generated": 1, "skipped": 1, "errors": 0}) + self.assertEqual([r["path"] for r in job.outcomes["skipped"]], ["a.png"]) + + async def test_rebuild_never_touches_originals(self): + for name in ("a", "b"): + make_png(self.output_dir / f"{name}.png", size=(400, 400)) + _, body = await self._start(size=128) + await self._finished(body["job"]["id"]) + originals = { + name: hashlib.sha256( + (self.output_dir / f"{name}.png").read_bytes() + ).hexdigest() + for name in ("a", "b") + } + + _, body = await self._start(mode="rebuild", size=256, quality="high") + job = await self._finished(body["job"]["id"]) + + self.assertEqual(job.counters, {"generated": 2, "skipped": 0, "errors": 0}) + for name in ("a", "b"): + digest = hashlib.sha256( + (self.output_dir / f"{name}.png").read_bytes() + ).hexdigest() + self.assertEqual(digest, originals[name], name) + with Image.open(self._thumb_path(name)) as thumb: + self.assertEqual(thumb.size, (256, 256)) + + async def test_rebuild_keeps_old_preview_when_generation_fails(self): + source = make_png(self.output_dir / "a.png", size=(64, 64)) + _, body = await self._start() + await self._finished(body["job"]["id"]) + preview = self._thumb_path("a").read_bytes() + source.write_bytes(b"\x89PNG broken now") + later = self._thumb_path("a").stat().st_mtime + 10 + os.utime(self._thumb_path("a"), (later, later)) # still fresh + + _, body = await self._start(mode="rebuild") + job = await self._finished(body["job"]["id"]) + + self.assertEqual(job.counters["errors"], 1) + self.assertEqual(self._thumb_path("a").read_bytes(), preview) + leftovers = [p.name for p in self._thumb_path("a").parent.iterdir()] + self.assertEqual(leftovers, ["a_thumb.png"]) # no temp file left + + async def test_rebuild_keeps_an_old_video_preview_instead_of_a_placeholder(self): + clip = self.output_dir / "clip.mp4" + clip.write_bytes(b"\x00") + preview = make_png(self.output_dir / "thumbnails" / "clip_thumb.jpg") + old = preview.read_bytes() + with ( + patch.dict(sys.modules, {"cv2": None}), + patch("subprocess.run", side_effect=FileNotFoundError("no ffmpeg")), + ): + _, body = await self._start(mode="rebuild") + job = await self._finished(body["job"]["id"]) + self.assertEqual(job.counters["errors"], 1) + self.assertEqual(preview.read_bytes(), old) + + async def test_quality_sets_the_encoder_not_the_edge(self): + noise = bytes( + (i * 7919 + (i >> 3) * 104729) % 251 for i in range(600 * 600 * 3) + ) + with Image.frombytes("RGB", (600, 600), noise) as img: + img.save(self.output_dir / "photo.jpg", "JPEG", quality=95) + sizes = {} + for quality in ("low", "high"): + _, body = await self._start(mode="rebuild", size=300, quality=quality) + await self._finished(body["job"]["id"]) + thumb = self._thumb_path("photo", ".jpg") + with Image.open(thumb) as image: + self.assertEqual(image.size, (300, 300)) + sizes[quality] = thumb.stat().st_size + self.assertLess(sizes["low"], sizes["high"]) + + async def test_no_output_folder_fails_the_job(self): + self.api._get_all_output_dirs = lambda: [] + _, body = await self._start() + job = await self._finished(body["job"]["id"]) + self.assertEqual(job.state, "failed") + self.assertEqual(job.status_line, "ComfyUI output directory not found") + + async def test_cancel_stops_at_file_boundary_keeping_written_previews(self): + for i in range(6): + make_png(self.output_dir / f"p{i}.png") + original = PromptManagerConfig.WORKER_THREADS + PromptManagerConfig.WORKER_THREADS = 1 + self.addCleanup(setattr, PromptManagerConfig, "WORKER_THREADS", original) + written = [] + real_write = images_module._write_image_thumbnail + job_ids = [] + + def counting(src, dst, *args, **kwargs): + real_write(src, dst, *args, **kwargs) + written.append(dst) + if len(written) == 2: + self.assertTrue(self.api._jobs.cancel(job_ids[0])) + + with patch.object(images_module, "_write_image_thumbnail", counting): + _, body = await self._start() + job_ids.append(body["job"]["id"]) + job = await self._finished(body["job"]["id"]) + + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.counters["generated"], 2) + self.assertEqual(len(written), 2) + for dst in written: + self.assertTrue(dst.is_file()) + made = sorted((self.output_dir / "thumbnails").glob("*_thumb.png")) + self.assertEqual(made, sorted(written)) + self.assertTrue(any("Thumbnails cancelled" in line for line in job.log)) + + async def test_a_second_thumbnail_job_is_refused_naming_the_first(self): + make_png(self.output_dir / "a.png") + entered, release = self._gate_generation() + _, first = await self._start() + await self._wait(entered) + + status, body = await self._start(mode="rebuild") + + self.assertEqual((status, body["code"]), (409, "busy")) + self.assertEqual(body["running"]["id"], first["job"]["id"]) + release.set() + await self._finished(first["job"]["id"]) + + async def test_thumbnail_job_refused_while_scan_runs(self): + make_png(self.output_dir / "a.png") + self._run_scan() + + status, body = await self._start() + + self.assertEqual((status, body["code"]), (409, "busy")) + self.assertEqual(body["running"]["title"], "Scan for metadata") + self.assertFalse(self._thumb_path("a").exists()) + + async def test_scan_refused_while_thumbnail_job_runs(self): + make_png(self.output_dir / "a.png") + entered, release = self._gate_generation() + _, started = await self._start() + await self._wait(entered) + + resp = await self.client.post("/prompt_manager/jobs/scan") + body = await resp.json() + + self.assertEqual((resp.status, body["code"]), (409, "busy")) + self.assertEqual(body["running"]["title"], "Thumbnail generation") + release.set() + await self._finished(started["job"]["id"]) + + async def test_clear_refused_while_thumbnail_job_runs(self): + make_png(self.output_dir / "a.png") + kept = make_png(self.output_dir / "thumbnails" / "old_thumb.png") + entered, release = self._gate_generation() + _, started = await self._start() + await self._wait(entered) + + resp = await self.client.post("/prompt_manager/images/clear-thumbnails") + + self.assertEqual(resp.status, 409) + self.assertEqual( + await resp.json(), + {"success": False, "error": "Thumbnail generation already running"}, + ) + self.assertTrue(kept.exists()) + release.set() + await self._finished(started["job"]["id"]) + + +# The legacy stream's frames and their keys, as the gallery page reads them. +LEGACY_KEYS = { + "status": {"phase", "message"}, + "start": { + "total_images", + "skipped", + "remaining", + "phase", + "image_count", + "video_count", + "message", + }, + "progress": { + "processed", + "total_images", + "generated", + "skipped", + "error_count", + "percentage", + "rate", + "eta", + "elapsed", + "current_file", + "file_type", + "action", + }, + "complete": { + "count", + "skipped", + "total_images", + "remaining", + "errors", + "error_count", + "elapsed_time", + "processing_rate", + "message", + }, +} + + +class TestLegacyStreamOnTheJob(ThumbnailJobCase): + async def test_legacy_progress_events_unchanged_on_success(self): + make_png(self.output_dir / "a.png", size=(400, 400)) + make_png(self.output_dir / "b.png") + + resp = await self.client.post(f"{self.PROGRESS_URL}?quality=low") + events = parse_sse(await resp.text()) + + self.assertEqual( + [name for name, _ in events], + ["status", "start", "progress", "progress", "complete"], + ) + for name, data in events: + self.assertEqual(set(data), LEGACY_KEYS[name], name) + complete = events[-1][1] + self.assertEqual((complete["count"], complete["total_images"]), (2, 2)) + self.assertEqual( + events[0][1], + {"phase": "scanning", "message": "Scanning for images and videos..."}, + ) + with Image.open(self._thumb_path("a")) as thumb: + self.assertEqual(max(thumb.size), 150) # the legacy preset still sizes + [job] = self.api._jobs.list() + self.assertEqual( + (job.kind, job.title, job.state, job.cancellable), + ("thumbnails", "Thumbnail generation", "finished", True), + ) + self.assertEqual(job.counters["generated"], 2) + + async def test_legacy_file_errors_still_stream(self): + make_png(self.output_dir / "good.png") + (self.output_dir / "bad.png").write_bytes(b"\x89PNG nope") + + resp = await self.client.post(self.PROGRESS_URL) + events = parse_sse(await resp.text()) + + self.assertEqual( + [d for e, d in events if e == "file_error"], + [{"file": "bad.png", "error": "Not a valid image file"}], + ) + + async def test_thumbnail_job_survives_client_disconnect(self): + for i in range(3): + make_png(self.output_dir / f"p{i}.png") + entered, release = self._gate_generation() + + resp = await self.client.post(self.PROGRESS_URL) + self.assertEqual(resp.status, 200) + await self._wait(entered) + job = self.api._jobs.active("thumbnails") + resp.close() # the gallery tab went away + release.set() + await asyncio.wait_for(job.task, 10) + + self.assertEqual(job.state, "finished") + self.assertEqual(job.counters["generated"], 3) + for i in range(3): + self.assertTrue(self._thumb_path(f"p{i}").is_file()) + + async def test_legacy_stream_attaches_to_a_running_job(self): + make_png(self.output_dir / "a.png") + entered, release = self._gate_generation() + _, started = await self._start(mode="missing", size=256) + await self._wait(entered) + + attached = asyncio.ensure_future(self.client.post(self.PROGRESS_URL)) + await asyncio.sleep(0) + job = self.api._jobs.get(started["job"]["id"]) + while job.subscriber_count == 0: + await asyncio.sleep(0.01) + release.set() + events = parse_sse(await (await attached).text()) + + self.assertEqual(events[-1][0], "complete") + # One job: the stream followed the running one with its own options. + [listed] = self.api._jobs.list() + self.assertIs(listed, job) + self.assertEqual(job.scope["size"], 256) + + async def test_cancel_from_activity_ends_the_legacy_stream_with_an_error(self): + for i in range(2): + make_png(self.output_dir / f"p{i}.png") + entered, release = self._gate_generation() + stream = asyncio.ensure_future(self.client.post(self.PROGRESS_URL)) + await self._wait(entered) + job = self.api._jobs.active("thumbnails") + self.assertTrue(self.api._jobs.cancel(job.id)) + release.set() + + events = parse_sse(await (await stream).text()) + + self.assertEqual(events[-1][0], "error") + self.assertEqual(events[-1][1]["error"], "Thumbnail generation was cancelled") + await asyncio.wait_for(job.task, 10) + self.assertEqual(job.state, "cancelled") + + +if __name__ == "__main__": + unittest.main() diff --git a/web/js/gallery.js b/web/js/gallery.js index cad3c2ca..3aa7d3ef 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -1911,6 +1911,24 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; } } + /** + * Stop the running thumbnail job (#249 System S5), whichever page started it: there is one at a time. + * Resolves whether a cancellation was asked for; a held legacy operation cannot be cancelled. + */ + async cancelThumbnailJob() { + try { + const data = await this.api.get('/prompt_manager/jobs'); + const job = (Array.isArray(data?.jobs) ? data.jobs : []).find(item => item?.kind === 'thumbnails' + && item.state === 'running' && item.cancellable !== false && typeof item.id === 'string'); + if (!job) return false; + await this.api.post(`/prompt_manager/jobs/${encodeURIComponent(job.id)}/cancel`, {}); + return true; + } catch (error) { + console.warn('The thumbnail job could not be cancelled:', error); + return false; + } + } + async generateThumbnailsWithProgress(options) { const progressBar = document.getElementById('progressBar'); const progressText = document.getElementById('progressText'); @@ -1933,11 +1951,13 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; let cancelled = false; let resultData = null; - // Handle cancellation + // Handle cancellation. The work is an Activity job that outlives this stream, so closing it no + // longer stops anything: ask the server to stop the job at its next file boundary too. const cancelHandler = () => { cancelled = true; eventSource.close(); cancelBtn.classList.add('hidden'); + this.cancelThumbnailJob(); reject(new Error('Thumbnail generation cancelled by user')); }; cancelBtn.addEventListener('click', cancelHandler);