From 554a213936d610b99a4c90878f3d4a98ffda4d62 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Fri, 16 Oct 2026 04:13:59 -0700 Subject: [PATCH] fix: count pending thumbnails accurately across gallery roots --- py/api/images.py | 85 +++++++++++++++++++++----- tests/js/gallery-thumbnails.test.js | 90 ++++++++++++++++++++++++++++ tests/test_api_images.py | 93 +++++++++++++++++++++++++++-- web/js/gallery.js | 67 +++++++++------------ 4 files changed, 277 insertions(+), 58 deletions(-) create mode 100644 tests/js/gallery-thumbnails.test.js diff --git a/py/api/images.py b/py/api/images.py index d03e27d..cf0a04c 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -214,9 +214,9 @@ def _thumbnail_complete_payload(total, stats, start, remaining=0, processed=None def _thumbnail_is_fresh(src, dst): - """True when ``dst`` exists and is newer than ``src`` (blocking stat).""" + """True when ``dst`` exists and is at least as new as ``src`` (blocking stat).""" try: - return dst.stat().st_mtime > src.stat().st_mtime + return dst.stat().st_mtime >= src.stat().st_mtime except OSError: return False @@ -382,6 +382,10 @@ class ImageRoutesMixin: async def delete_image_route(request): return await self.delete_image(request) + @routes.get("/prompt_manager/images/thumbnails/status") + async def thumbnail_status_route(request): + return await self.get_thumbnail_status(request) + @routes.post("/prompt_manager/images/generate-thumbnails") async def generate_thumbnails_route(request): return await self.generate_thumbnails(request) @@ -815,19 +819,15 @@ class ImageRoutesMixin: return error thumbnail_size = self._thumbnail_size(data.get("quality", "medium")) - output_dir = self._find_comfyui_output_dir() - if not output_dir: + output_dirs = self._get_all_output_dirs() + if not output_dirs: return web.json_response( {"success": False, "error": "ComfyUI output directory not found"}, status=404, ) - output_path = Path(output_dir) result = await self._run_in_executor( - self._generate_thumbnails_sync, - output_path, - output_path / "thumbnails", - thumbnail_size, + self._generate_gallery_thumbnails_sync, output_dirs, thumbnail_size ) self.invalidate_gallery_cache() return web.json_response(result) @@ -842,6 +842,22 @@ class ImageRoutesMixin: """Blocking thumbnail generation loop (run in executor).""" start = _time.monotonic() plan = self._plan_thumbnails(output_path, thumbnails_dir) + payload = self._execute_thumbnail_plan(plan, thumbnail_size, start) + payload["thumbnails_path"] = _public_path(thumbnails_dir) + return payload + + def _generate_gallery_thumbnails_sync( + self, output_dirs: list[Path], thumbnail_size: tuple[int, int] + ) -> dict: + start = _time.monotonic() + plan = self._plan_gallery_thumbnails(output_dirs) + payload = self._execute_thumbnail_plan(plan, thumbnail_size, start) + payload["thumbnails_path"] = _public_path(output_dirs[0] / "thumbnails") + return payload + + def _execute_thumbnail_plan( + self, plan: dict, thumbnail_size: tuple[int, int], start: float + ) -> dict: targets = plan["targets"] stats = _new_thumbnail_stats() stats["skipped"] = plan["skipped"] @@ -855,10 +871,50 @@ class ImageRoutesMixin: processed=len(targets), ) payload["success"] = True - payload["thumbnails_path"] = _public_path(thumbnails_dir) 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} + for root in output_dirs: + root = Path(root) + plan = self._plan_thumbnails(root, root / "thumbnails") + available = max(0, MAX_THUMBNAIL_FILES - len(combined["targets"])) + combined["targets"].extend(plan["targets"][:available]) + combined["skipped"] += plan["skipped"] + combined["remaining"] += plan["remaining"] + max( + 0, len(plan["targets"]) - available + ) + return combined + + async def get_thumbnail_status(self, request: web.Request) -> web.Response: + """Count pending thumbnails off-loop without generating files.""" + try: + output_dirs = self._get_all_output_dirs() + plan = await self._run_in_executor( + self._plan_gallery_thumbnails, output_dirs + ) + pending = len(plan["targets"]) + plan["remaining"] + collection = json.dumps(sorted(str(_canonical(p)) for p in output_dirs)) + return web.json_response( + { + "success": True, + "total_images": pending + plan["skipped"], + "pending": pending, + "ready": plan["skipped"], + "collection_id": hashlib.sha256(collection.encode()).hexdigest()[ + :16 + ], + }, + headers={"Cache-Control": "no-store"}, + ) + except Exception as exc: + self.logger.exception("Thumbnail status failed") + return web.json_response( + {"success": False, "error": _safe_error(exc)}, status=500 + ) + def _plan_thumbnails(self, output_path, thumbnails_dir): """Decide which media files need a thumbnail (blocking walk). @@ -923,7 +979,7 @@ class ImageRoutesMixin: try: if not src.is_file(): return {**result, "action": "error", "error": "File no longer exists"} - if dst.exists() and dst.stat().st_mtime > src.stat().st_mtime: + if _thumbnail_is_fresh(src, dst): return {**result, "action": "skipped"} dst.parent.mkdir(parents=True, exist_ok=True) if is_video: @@ -957,21 +1013,20 @@ class ImageRoutesMixin: ) await response.prepare(request) try: - output_dir = self._find_comfyui_output_dir() - if not output_dir: + output_dirs = self._get_all_output_dirs() + if not output_dirs: await self._emit_sse( response, "error", {"error": "ComfyUI output directory not found"} ) return response - output_path = Path(output_dir) await self._emit_sse( response, "status", {"phase": "scanning", "message": "Scanning for images and videos..."}, ) plan = await self._run_in_executor( - self._plan_thumbnails, output_path, output_path / "thumbnails" + self._plan_gallery_thumbnails, output_dirs ) await self._stream_thumbnails( response, diff --git a/tests/js/gallery-thumbnails.test.js b/tests/js/gallery-thumbnails.test.js new file mode 100644 index 0000000..535e945 --- /dev/null +++ b/tests/js/gallery-thumbnails.test.js @@ -0,0 +1,90 @@ +const test = require('node:test'); +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const path = require('node:path'); +const vm = require('node:vm'); + +function fixture(status, settings = {}) { + const timers = []; + const prompts = []; + const requests = []; + const buttons = {}; + const storage = new Map(); + const document = { + addEventListener() {}, + createElement: () => ({ addEventListener() {} }), + body: { appendChild() {}, removeChild() {} }, + getElementById: id => buttons[id] ||= { addEventListener(type, handler) { this[type] = handler; } }, + }; + const context = vm.createContext({ document, window: {}, console: { debug() {}, warn() {} }, + setTimeout: callback => timers.push(callback), + localStorage: { setItem: (key, value) => storage.set(key, value) }, + ApiClient: { fetchUrl: async () => ({ json: async () => ({ success: true, + total: 10000, images: Array.from({ length: 50 }, () => ({ thumbnail_url: null })), + }) }) }, + }); + const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context); + const manager = Object.create(Manager.prototype); + Object.assign(manager, { getSettings: () => ({ checkThumbnailsAtStartup: true, ...settings }), + api: { get: async url => { requests.push(url); return status; } }, + }); + const originalPrompt = manager.showThumbnailGenerationPrompt.bind(manager); + manager.showThumbnailGenerationPrompt = (...args) => prompts.push(args); + async function check() { + await manager.checkThumbnailsAtStartup(); + await timers.shift()(); + } + return { manager, check, prompts, requests, settings, buttons, storage, originalPrompt }; +} + +test('startup uses exact collection counts instead of extrapolating the newest 50 images', async () => { + const f = fixture({ success: true, total_images: 10000, pending: 50, ready: 9950, collection_id: 'a' }); + await f.check(); + assert.deepEqual(f.requests, ['/prompt_manager/images/thumbnails/status']); + assert.equal(f.prompts.length, 0); +}); + +test('startup prompt reports the actual number needing generation', async () => { + const f = fixture({ success: true, total_images: 10000, pending: 1200, collection_id: 'a' }); + await f.check(); + assert.equal(f.prompts.length, 1); + assert.deepEqual(f.prompts[0].slice(0, 2), [10000, 1200]); +}); + +test('a complete collection never prompts again', async () => { + const f = fixture({ success: true, total_images: 10000, pending: 0, collection_id: 'a' }); + await f.check(); + await f.check(); + assert.equal(f.prompts.length, 0); +}); + +test('Not Now persists the collection state and suppresses the unchanged prompt', async () => { + const status = { success: true, total_images: 100, pending: 30, collection_id: 'a' }; + const f = fixture(status); + await f.check(); + f.originalPrompt(...f.prompts[0]); + f.buttons.thumbnailPromptNo.click(); + Object.assign(f.settings, JSON.parse(f.storage.get('gallerySettings'))); + f.prompts.length = 0; + await f.check(); + assert.equal(f.prompts.length, 0); + status.pending = 40; + await f.check(); + assert.equal(f.prompts.length, 1); +}); + +test('disabled checks do not request collection status', async () => { + const f = fixture({}, { checkThumbnailsAtStartup: false }); + await f.check(); + assert.equal(f.requests.length, 0); + assert.equal(f.prompts.length, 0); +}); + +test('saving gallery settings retains the remembered thumbnail dismissal', () => { + const f = fixture({}, { thumbnailPromptDismissed: 'remembered-state' }); + f.manager.showNotification = () => {}; + f.manager.hideSettings = () => {}; + f.manager.applySettings = () => {}; + f.manager.saveSettings(); + assert.equal(JSON.parse(f.storage.get('gallerySettings')).thumbnailPromptDismissed, 'remembered-state'); +}); diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 9da9829..be6a3e4 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -486,7 +486,7 @@ class TestThumbnailProgressStream(ThumbnailTestCase): # PIL work ran in the executor in batches (never on the loop), plus the scan. self.assertGreaterEqual(calls.count("_generate_batch"), 1) self.assertNotIn("_generate_one", calls) - self.assertIn("_plan_thumbnails", calls) + self.assertIn("_plan_gallery_thumbnails", calls) async def test_capped_run_reports_remaining_and_the_next_run_finishes(self): for name in ("a", "b", "c"): @@ -573,7 +573,7 @@ class TestThumbnailProgressStream(ThumbnailTestCase): self.assertEqual(events[-1][1]["total_images"], 0) async def test_missing_output_dir_sends_error_event(self): - self.api._find_comfyui_output_dir = lambda: None + self.api._get_all_output_dirs = lambda: [] resp = await self.client.request("POST", self.PROGRESS_URL) @@ -616,7 +616,7 @@ class TestGenerateThumbnailsPost(ThumbnailTestCase): self.assertEqual(data2["count"], 0) async def test_missing_output_dir_is_404(self): - self.api._find_comfyui_output_dir = lambda: None + self.api._get_all_output_dirs = lambda: [] resp = await self.client.request( "POST", "/prompt_manager/images/generate-thumbnails", json={} @@ -1420,7 +1420,7 @@ class TestThumbnailUnits(ImageRouteCoverageCase): self.assertEqual(events[-1][0], "error") async def test_post_generate_reports_internal_failure(self): - self.api._find_comfyui_output_dir = _raise + self.api._get_all_output_dirs = _raise status, _ = await self._json( "POST", "/prompt_manager/images/generate-thumbnails", json={} ) @@ -1819,3 +1819,88 @@ class TestOutputScanHelpers(ImageAPITestCase): 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.""" + + STATUS_URL = "/prompt_manager/images/thumbnails/status" + + async def test_status_reports_exact_pending_count_without_generating(self): + for name in ("ready", "missing", "outdated"): + make_png(self.output_dir / f"{name}.png") + thumbs = self.output_dir / "thumbnails" + ready = make_png(thumbs / "ready_thumb.png") + outdated = make_png(thumbs / "outdated_thumb.png") + source_time = (self.output_dir / "ready.png").stat().st_mtime + os.utime(ready, (source_time + 10, source_time + 10)) + os.utime(outdated, (source_time - 10, source_time - 10)) + calls = self._spy_executor() + + response = await self.client.get(self.STATUS_URL) + + self.assertEqual(response.status, 200) + data = await response.json() + self.assertTrue(data["success"]) + self.assertEqual(data["total_images"], 3) + self.assertEqual(data["pending"], 2) + self.assertEqual(data["ready"], 1) + self.assertTrue(data["collection_id"]) + self.assertIn("_plan_gallery_thumbnails", calls) + self.assertFalse((thumbs / "missing_thumb.png").exists()) + + async def test_both_generation_routes_cover_all_gallery_roots(self): + second = self.comfy_dir / "another-output" + make_png(self.output_dir / "a.png") + make_png(second / "b.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + for route in ( + "/prompt_manager/images/generate-thumbnails", + self.PROGRESS_URL, + ): + with self.subTest(route=route): + for root in (self.output_dir, second): + for thumb in (root / "thumbnails").glob("*_thumb.png"): + thumb.unlink() + response = await self.client.post(route, json={}) + if route == self.PROGRESS_URL: + data = dict(parse_sse(await response.text()))["complete"] + else: + data = await response.json() + self.assertEqual(data["count"], 2) + self.assertTrue((second / "thumbnails" / "b_thumb.png").exists()) + status = await (await self.client.get(self.STATUS_URL)).json() + self.assertEqual(status["pending"], 0) + self.assertEqual(status["total_images"], 2) + + async def test_pending_count_includes_files_beyond_the_generation_cap(self): + for name in ("a", "b", "c"): + make_png(self.output_dir / f"{name}.png") + with patch.object(images_module, "MAX_THUMBNAIL_FILES", 1): + response = await self.client.get(self.STATUS_URL) + self.assertEqual(response.status, 200) + data = await response.json() + self.assertEqual(data["pending"], 3) + + async def test_generation_cap_is_shared_across_roots(self): + second = self.comfy_dir / "another-output" + make_png(self.output_dir / "a.png") + make_png(second / "b.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + with patch.object(images_module, "MAX_THUMBNAIL_FILES", 1): + response = await self.client.post(self.PROGRESS_URL) + data = dict(parse_sse(await response.text()))["complete"] + self.assertEqual(data["count"], 1) + self.assertEqual(data["remaining"], 1) + self.assertEqual(data["total_images"], 2) + + async def test_equal_timestamps_do_not_regenerate_existing_thumbnails(self): + source = make_png(self.output_dir / "ready.png") + thumb = make_png(self.output_dir / "thumbnails" / "ready_thumb.png") + timestamp = int(source.stat().st_mtime) + os.utime(source, (timestamp, timestamp)) + os.utime(thumb, (timestamp, timestamp)) + response = await self.client.post(self.PROGRESS_URL) + data = dict(parse_sse(await response.text()))["complete"] + self.assertEqual(data["count"], 0) + self.assertEqual(data["skipped"], 1) diff --git a/web/js/gallery.js b/web/js/gallery.js index 5dcea0e..836f54a 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -1747,6 +1747,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; debugMode: document.getElementById('debugModeToggle').checked, apiTimeout: parseInt(document.getElementById('apiTimeoutInput').value), thumbnailsGenerated: this.getSettings().thumbnailsGenerated, // Preserve this + thumbnailPromptDismissed: this.getSettings().thumbnailPromptDismissed, checkThumbnailsAtStartup: document.getElementById('checkThumbnailsAtStartup').checked, // Video settings videoAutoplay: document.getElementById('videoAutoplayToggle').checked, @@ -2503,37 +2504,16 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; } try { - // Check a sample of images to see if any are missing thumbnails - const response = await ApiClient.fetchUrl('/prompt_manager/images/output?limit=50&offset=0'); - const data = await response.json(); - - console.debug('Checking images for missing thumbnails:', data); - - if (data.success && data.images && data.images.length > 0) { - // Count images missing thumbnails - let missingThumbnails = 0; - for (const image of data.images) { - if (!image.thumbnail_url) { - missingThumbnails++; - } - } - - console.debug(`Found ${missingThumbnails} images without thumbnails in sample of ${data.images.length}`); - - // If more than 10% of sampled images are missing thumbnails, show prompt - const missingPercentage = (missingThumbnails / data.images.length) * 100; - if (missingThumbnails > 0 && missingPercentage > 10) { - // Estimate total missing based on sample - const estimatedMissing = Math.round((missingThumbnails / data.images.length) * data.total); - console.debug(`Estimated ${estimatedMissing} total images missing thumbnails (${missingPercentage.toFixed(1)}% of sample)`); - this.showThumbnailGenerationPrompt(data.total, estimatedMissing); - } else if (missingThumbnails > 0) { - console.debug(`Only ${missingPercentage.toFixed(1)}% of images missing thumbnails, not showing prompt`); - } else { - console.debug('All sampled images have thumbnails'); - } - } else { - console.debug('No images found or failed to fetch images'); + const data = await this.api.get('/prompt_manager/images/thumbnails/status'); + const total = data.total_images; + const pending = data.pending; + if (!data.success || !Number.isSafeInteger(total) || + !Number.isSafeInteger(pending) || total <= 0 || + pending <= 0 || pending > total) return; + const signature = JSON.stringify([data.collection_id, total, pending]); + if (settings.thumbnailPromptDismissed === signature) return; + if (pending / total > 0.1) { + this.showThumbnailGenerationPrompt(total, pending, signature); } } catch (error) { console.debug('Could not check for missing thumbnails at startup:', error); @@ -2541,7 +2521,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; }, 1000); // Wait 1 second for page to settle } - showThumbnailGenerationPrompt(imageCount, missingCount = imageCount) { + showThumbnailGenerationPrompt(imageCount, missingCount = imageCount, signature = null) { const modal = document.createElement('div'); modal.className = 'fixed inset-0 bg-black/50 flex items-center justify-center z-50'; modal.innerHTML = ` @@ -2552,10 +2532,10 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; Generate Thumbnails?

- Found ${missingCount} images without thumbnails (${imageCount} total images). Would you like to generate thumbnails for faster loading? + ${missingCount} images need new or updated thumbnails (${imageCount} total images). Generate them for faster loading?

- 💡 Thumbnails speed up gallery loading significantly. You can change this setting in Gallery Settings. + Only missing or outdated thumbnails will be generated. Existing current thumbnails will be kept. You can change startup checks in Gallery Settings.