From b3e031790dcd362ec1b24abff942a21660292807 Mon Sep 17 00:00:00 2001 From: Simon Cordt Date: Wed, 24 Sep 2025 10:03:29 +0200 Subject: [PATCH] Add image preview API to Persistent Queue --- server/queue_manager.py | 120 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) diff --git a/server/queue_manager.py b/server/queue_manager.py index b9a3344..86a24f0 100644 --- a/server/queue_manager.py +++ b/server/queue_manager.py @@ -9,6 +9,12 @@ from typing import Optional, Any, Dict, Tuple, List from aiohttp import web from .database import QueueDatabase +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from io import BytesIO +import os +import hashlib +import folder_paths class PersistentQueueManager: def __init__(self): @@ -285,6 +291,7 @@ class PersistentQueueManager: app.add_routes([ web.get('/api/pqueue', self._api_get_pqueue), web.get('/api/pqueue/history', self._api_get_history), + web.get('/api/pqueue/preview', self._api_preview_image), web.post('/api/pqueue/pause', self._api_pause), web.post('/api/pqueue/resume', self._api_resume), web.post('/api/pqueue/reorder', self._api_reorder), @@ -294,4 +301,117 @@ class PersistentQueueManager: except Exception as e: logging.debug(f"PersistentQueue add routes failed: {e}") + async def _api_preview_image(self, request: web.Request) -> web.Response: + try: + filename = request.rel_url.query.get('filename') + subfolder = request.rel_url.query.get('subfolder', '') + folder_type = request.rel_url.query.get('type', 'output') + pid = request.rel_url.query.get('pid') + preview_q = request.rel_url.query.get('preview', 'webp;50') + preview_info = preview_q.split(';') + image_format = preview_info[0] if preview_info else 'webp' + if image_format not in ['webp', 'jpeg', 'png']: + image_format = 'webp' + quality = 90 + if len(preview_info) > 1 and preview_info[-1].isdigit(): + quality = int(preview_info[-1]) + + # Resolve file path similar to /view + base_dir = folder_paths.get_directory_by_type(folder_type) + if base_dir is None: + return web.Response(status=400) + output_dir = base_dir + if subfolder: + full_output_dir = os.path.join(base_dir, subfolder) + if os.path.commonpath((os.path.abspath(full_output_dir), base_dir)) != base_dir: + return web.Response(status=403) + output_dir = full_output_dir + filename = os.path.basename(filename or '') + file = os.path.join(output_dir, filename) + + if not os.path.isfile(file): + return web.Response(status=404) + + # Build a cache path based on file + params + workflow hash + temp_dir = folder_paths.get_temp_directory() + cache_dir = os.path.join(temp_dir, 'preview_cache') + os.makedirs(cache_dir, exist_ok=True) + + workflow_json = None + if pid: + # Prefer workflow from history if available, else from queue + # Try job_history first + try: + # Direct DB access to last history for pid + with self.db._get_conn() as conn: + cur = conn.execute('SELECT workflow FROM job_history WHERE prompt_id = ? ORDER BY id DESC LIMIT 1', (pid,)) + row = cur.fetchone() + if row and row['workflow']: + workflow_json = row['workflow'] + except Exception: + workflow_json = None + if workflow_json is None: + job = self.db.get_job(pid) + if job: + workflow_json = job.get('workflow') + + # Include workflow_json in cache key so regenerated when workflow changes + wf_hash = hashlib.sha256((workflow_json or '').encode('utf-8')).hexdigest()[:16] + stat = os.stat(file) + cache_key = f"{os.path.abspath(file)}|{stat.st_mtime}|{stat.st_size}|{image_format}|{quality}|{wf_hash}" + cache_name = hashlib.sha256(cache_key.encode('utf-8')).hexdigest() + f'.{image_format}' + cache_path = os.path.join(cache_dir, cache_name) + + if os.path.exists(cache_path): + return web.FileResponse(cache_path, headers={"Content-Disposition": f"filename=\"{filename}\""}) + + with Image.open(file) as img: + save_kwargs: Dict[str, Any] = {"format": image_format} + if image_format in ['webp', 'jpeg']: + save_kwargs['quality'] = quality + + # Embed metadata if WEBP or PNG + if image_format == 'webp': + try: + exif = img.getexif() + # Copy existing PNG text into WEBP EXIF + if hasattr(img, 'text'): + for k in img.text: + val = img.text[k] + # Preserve both prompt and other keys + if k == 'prompt': + exif[0x0110] = "prompt:{}".format(val) + else: + tag = 0x010F + try: + exif[tag] = f"{k}:{val}" + except Exception: + pass + # Also embed provided workflow explicitly if given + if workflow_json: + try: + exif[0x010E] = "workflow:{}".format(workflow_json) # 0x010E = ImageDescription + except Exception: + pass + save_kwargs['exif'] = exif + except Exception: + pass + elif image_format == 'png': + try: + pnginfo = PngInfo() + if hasattr(img, 'text'): + for k in img.text: + pnginfo.add_text(k, img.text[k]) + if workflow_json: + pnginfo.add_text('workflow', workflow_json) + save_kwargs['pnginfo'] = pnginfo + except Exception: + pass + + img.save(cache_path, **save_kwargs) + + return web.FileResponse(cache_path, headers={"Content-Disposition": f"filename=\"{filename}\""}) + except Exception: + return web.Response(status=500) + queue_manager = PersistentQueueManager() \ No newline at end of file