nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(system): run thumbnail generation as an Activity job

Browse Source 
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.
This commit is contained in:
Vito Sansevero
2026-10-23 13:31:19 -07:00
parent 69de20a427
commit d5bcae301b
11 changed files with 1258 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TESTING.md
+25
View File
@@ -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
py/api/__init__.py
+3
View File
@@ -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,
py/api/images.py
+125 -172
View File
@@ -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
py/api/job_routes.py
+6 -1
View File
@@ -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)
py/api/jobs.py
+29
View File
@@ -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.
py/api/thumbnail_jobs.py
+448
View File
@@ -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
tests/js/gallery-thumbnails.test.js
+31
View File
@@ -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, []);
});
tests/test_api_images.py
+7 -25
View File
@@ -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."""
tests/test_jobs.py
+22
View File
@@ -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")
tests/test_thumbnail_jobs.py
+541
View File
@@ -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()
web/js/gallery.js
+21 -1
View File
@@ -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);