diff --git a/__init__.py b/__init__.py index da092da..1f0efbb 100644 --- a/__init__.py +++ b/__init__.py @@ -9,8 +9,9 @@ WEB_DIRECTORY = "./web" # Install persistent queue hooks as soon as this extension is imported. try: - from .server.queue_manager import queue_manager - queue_manager.initialize() + from .server.manager import PersistentQueueManager + _queue_manager = PersistentQueueManager() + _queue_manager.initialize() except Exception as e: logging.debug(f"ComfyUI-Persistent-Queue initialization failed: {e}") diff --git a/server/queue_manager.py b/server/manager.py similarity index 72% rename from server/queue_manager.py rename to server/manager.py index 4c63562..74b842d 100644 --- a/server/queue_manager.py +++ b/server/manager.py @@ -3,7 +3,6 @@ import json import logging import time import heapq -import copy from typing import Optional, Any, Dict, Tuple, List, Callable, Set from aiohttp import web @@ -15,305 +14,11 @@ from io import BytesIO import os import hashlib import folder_paths -import json -from PIL import ImageDraw, ImageFont +from .thumbnail_service import ThumbnailService +from .queue_hook_manager import QueueHookManager +from .routes_helper import RoutesHelper -class ThumbnailService: - """Generates small, web-friendly thumbnails from ComfyUI output descriptors. - - Single responsibility: image IO and thumbnail encoding. - """ - - def __init__(self, max_size: int = 128, quality: int = 60): - self.max_size = max_size - self.quality = quality - - def generate_thumbnails_from_outputs(self, outputs: Optional[dict], *, workflow_json: Optional[str] = None, extras: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: - if not outputs: - return [] - images: List[Dict[str, Any]] = self._extract_image_descriptors(outputs) - thumbs: List[Dict[str, Any]] = [] - for idx, desc in enumerate(images[:4]): - thumb = self._encode_single_thumbnail(desc, idx, workflow_json=workflow_json, extras=extras) - if thumb is not None: - thumbs.append(thumb) - return thumbs - - def generate_placeholder_thumbnail(self, status: str, *, workflow_json: Optional[str] = None) -> Optional[Dict[str, Any]]: - """Create a placeholder WEBP thumbnail for failed/interrupted jobs. - - The placeholder encodes the workflow in EXIF so it can be restored via drag-and-drop. - """ - try: - size = int(self.max_size) - # Locate provided placeholder asset - placeholder_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'web', 'img', 'failed.png')) - use_draw_fallback = True - if os.path.isfile(placeholder_path): - try: - with Image.open(placeholder_path) as base_img: - # Resize to fit within max size, maintain aspect ratio - w, h = base_img.size - scale = min(size / max(1, w), size / max(1, h), 1.0) - new_size = (max(1, int(w * scale)), max(1, int(h * scale))) - if hasattr(Image, 'Resampling'): - resampling = Image.Resampling.LANCZOS - else: - resampling = Image.LANCZOS - img = base_img.convert('RGB').resize(new_size, resampling) - use_draw_fallback = False - except Exception: - use_draw_fallback = True - - if use_draw_fallback: - # Minimal drawn fallback if asset missing/unreadable - bg = (28, 29, 32) - color = (239, 68, 68) if (status or '').lower() == 'failed' else (234, 179, 8) - text = '×' if (status or '').lower() == 'failed' else '!' - img = Image.new('RGB', (size, size), bg) - draw = ImageDraw.Draw(img) - margin = 12 - rect = (margin, margin, size - margin, size - margin) - try: - draw.rounded_rectangle(rect, radius=10, fill=color) - except Exception: - draw.rectangle(rect, fill=color) - try: - font = ImageFont.load_default() - except Exception: - font = None - if font is not None: - try: - bbox = draw.textbbox((0, 0), text, font=font) - tw = bbox[2] - bbox[0] - th = bbox[3] - bbox[1] - except Exception: - tw, th = draw.textlength(text, font=font), 10 - tx = (size - tw) // 2 - ty = (size - th) // 2 - draw.text((tx, ty), text, font=font, fill=(255, 255, 255)) - - buf = BytesIO() - try: - pnginfo = PngInfo() - if workflow_json: - pnginfo.add_text('prompt', workflow_json) - # Save lossless PNG to preserve colors and avoid artifacts - img.save(buf, format='PNG', pnginfo=pnginfo, compress_level=4) - except Exception: - img.save(buf, format='PNG', compress_level=4) - data = buf.getvalue() - return { - 'idx': 0, - 'mime': 'image/png', - 'width': img.size[0], - 'height': img.size[1], - 'data': data, - } - except Exception: - return None - - def _extract_image_descriptors(self, outputs: dict) -> List[Dict[str, Any]]: - images: List[Dict[str, Any]] = [] - try: - for v in (outputs or {}).values(): - if isinstance(v, dict): - imgs = v.get('images') or (v.get('ui') or {}).get('images') or [] - if isinstance(imgs, list): - for i in imgs: - if isinstance(i, dict) and (i.get('filename') or i.get('name')): - images.append(i) - elif isinstance(v, list): - for i in v: - if isinstance(i, dict) and (i.get('filename') or i.get('name')): - images.append(i) - except Exception: - return [] - return images - - def _encode_single_thumbnail(self, desc: Dict[str, Any], idx: int, *, workflow_json: Optional[str], extras: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: - filename = desc.get('filename') or desc.get('name') - folder_type = desc.get('type') or 'output' - subfolder = desc.get('subfolder') or '' - - base_dir = folder_paths.get_directory_by_type(folder_type) - if base_dir is None: - return None - img_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 None - img_dir = full_output_dir - file_path = os.path.join(img_dir, os.path.basename(filename)) - if not os.path.isfile(file_path): - return None - - try: - with Image.open(file_path) as img: - # Resize preserving aspect ratio to fit within max_size - w, h = img.size - scale = min(self.max_size / max(1, w), self.max_size / max(1, h), 1.0) - new_size = (max(1, int(w * scale)), max(1, int(h * scale))) - if hasattr(Image, 'Resampling'): - resampling = Image.Resampling.LANCZOS - else: - resampling = Image.LANCZOS - thumb_img = img.convert('RGB').resize(new_size, resampling) - buf = BytesIO() - # Embed workflow metadata into WEBP EXIF so drag-and-drop works - try: - exif = thumb_img.getexif() - if workflow_json: - # Comfy expects 'prompt:' in 0x0110 for WEBP - exif[0x0110] = "prompt:{}".format(workflow_json) - if extras: - tag = 0x010F - for k, v in extras.items(): - try: - exif[tag] = f"{k}:{json.dumps(v)}" - except Exception: - pass - tag -= 1 - thumb_img.save(buf, format='WEBP', quality=self.quality, exif=exif) - except Exception: - thumb_img.save(buf, format='WEBP', quality=self.quality) - data = buf.getvalue() - return { - 'idx': idx, - 'mime': 'image/webp', - 'width': new_size[0], - 'height': new_size[1], - 'data': data, - } - except Exception: - return None - -class QueueHookManager: - """Manages installation/uninstallation of queue hooks. - - Responsible for wrapping prompt queue methods to add persistence and pause behavior. - """ - - def __init__(self, *, is_paused_fn: Callable[[], bool], on_job_started: Callable[[str], None], on_task_done: Callable[[Any, Any, Any], None], should_run_when_paused: Optional[Callable[[str], bool]] = None): - self._original_queue_get = None - self._original_task_done = None - self._installed = False - self._is_paused = is_paused_fn - self._on_job_started = on_job_started - self._on_task_done = on_task_done - self._should_run_when_paused = should_run_when_paused - - def install(self) -> None: - """Install hooks into execution.PromptQueue if not already installed.""" - if self._installed: - return - import execution - - if self._original_queue_get is None: - self._original_queue_get = execution.PromptQueue.get - - def get_wrapper(q_self, timeout=None): - # If paused, only allow items explicitly permitted by the manager (run-selected mode) - try: - if self._is_paused(): - allowed = False - if callable(self._should_run_when_paused): - try: - top = q_self.queue[0] if getattr(q_self, 'queue', None) else None - except Exception: - top = None - if top is not None: - try: - top_pid = str(top[1]) - allowed = bool(self._should_run_when_paused(top_pid)) - except Exception: - allowed = False - if not allowed: - time.sleep(0.1) - return None - except Exception: - # If anything goes wrong during checks, be safe and respect pause - if self._is_paused(): - time.sleep(0.1) - return None - - result = self._original_queue_get(q_self, timeout=timeout) - if result is not None: - try: - item, _item_id = result - prompt_id = item[1] - # If paused, ensure the popped item is allowed; otherwise, reinsert and yield None - if self._is_paused() and callable(self._should_run_when_paused): - try: - if not self._should_run_when_paused(str(prompt_id)): - try: - q_self.queue.append(item) - heapq.heapify(q_self.queue) - except Exception: - pass - time.sleep(0.05) - return None - except Exception: - time.sleep(0.05) - return None - self._on_job_started(prompt_id) - except Exception as e: - logging.debug(f"QueueHookManager get_wrapper failed: {e}") - return result - - execution.PromptQueue.get = get_wrapper - - if self._original_task_done is None: - self._original_task_done = execution.PromptQueue.task_done - - def task_done_wrapper(q_self, item_id, history_result, status): - try: - self._on_task_done((q_self, item_id, history_result, status)) - except Exception as e: - logging.debug(f"QueueHookManager task_done on_task_done failed: {e}") - self._original_task_done(q_self, item_id, history_result, status) - - execution.PromptQueue.task_done = task_done_wrapper - - self._installed = True - - def uninstall(self) -> None: - """Restore original methods.""" - if not self._installed: - return - import execution - if self._original_queue_get is not None: - execution.PromptQueue.get = self._original_queue_get - self._original_queue_get = None - if self._original_task_done is not None: - execution.PromptQueue.task_done = self._original_task_done - self._original_task_done = None - self._installed = False - -class RoutesHelper: - """Registers HTTP routes for the persistent queue API.""" - - def __init__(self, app: Any): - self.app = app - - def register(self, manager: "PersistentQueueManager") -> None: - routes = [ - web.get('/api/pqueue', manager._api_get_pqueue), - web.get('/api/pqueue/history', manager._api_get_history), - web.get('/api/pqueue/history/thumb/{history_id:\\d+}', manager._api_get_history_thumb), - web.get('/api/pqueue/preview', manager._api_preview_image), - web.post('/api/pqueue/pause', manager._api_pause), - web.post('/api/pqueue/resume', manager._api_resume), - web.post('/api/pqueue/reorder', manager._api_reorder), - web.patch('/api/pqueue/priority', manager._api_priority), - web.post('/api/pqueue/delete', manager._api_delete), - web.patch('/api/pqueue/rename', manager._api_rename), - web.post('/api/pqueue/run-selected', manager._api_run_selected), - web.post('/api/pqueue/skip-selected', manager._api_skip_selected), - ] - self.app.add_routes(routes) class PersistentQueueManager: """Coordinator for persistent queue persistence, API handlers, and hooks.""" @@ -1081,4 +786,4 @@ class PersistentQueueManager: except Exception: pass -queue_manager = PersistentQueueManager() \ No newline at end of file + diff --git a/server/queue_hook_manager.py b/server/queue_hook_manager.py new file mode 100644 index 0000000..5007142 --- /dev/null +++ b/server/queue_hook_manager.py @@ -0,0 +1,109 @@ +import time +import heapq +import logging +from typing import Optional, Any, Callable + + +class QueueHookManager: + """Manages installation/uninstallation of queue hooks. + + Responsible for wrapping prompt queue methods to add persistence and pause behavior. + """ + + def __init__(self, *, is_paused_fn: Callable[[], bool], on_job_started: Callable[[str], None], on_task_done: Callable[[Any, Any, Any], None], should_run_when_paused: Optional[Callable[[str], bool]] = None): + self._original_queue_get = None + self._original_task_done = None + self._installed = False + self._is_paused = is_paused_fn + self._on_job_started = on_job_started + self._on_task_done = on_task_done + self._should_run_when_paused = should_run_when_paused + + def install(self) -> None: + """Install hooks into execution.PromptQueue if not already installed.""" + if self._installed: + return + import execution + + if self._original_queue_get is None: + self._original_queue_get = execution.PromptQueue.get + + def get_wrapper(q_self, timeout=None): + # If paused, only allow items explicitly permitted by the manager (run-selected mode) + try: + if self._is_paused(): + allowed = False + if callable(self._should_run_when_paused): + try: + top = q_self.queue[0] if getattr(q_self, 'queue', None) else None + except Exception: + top = None + if top is not None: + try: + top_pid = str(top[1]) + allowed = bool(self._should_run_when_paused(top_pid)) + except Exception: + allowed = False + if not allowed: + time.sleep(0.1) + return None + except Exception: + # If anything goes wrong during checks, be safe and respect pause + if self._is_paused(): + time.sleep(0.1) + return None + + result = self._original_queue_get(q_self, timeout=timeout) + if result is not None: + try: + item, _item_id = result + prompt_id = item[1] + # If paused, ensure the popped item is allowed; otherwise, reinsert and yield None + if self._is_paused() and callable(self._should_run_when_paused): + try: + if not self._should_run_when_paused(str(prompt_id)): + try: + q_self.queue.append(item) + heapq.heapify(q_self.queue) + except Exception: + pass + time.sleep(0.05) + return None + except Exception: + time.sleep(0.05) + return None + self._on_job_started(prompt_id) + except Exception as e: + logging.debug(f"QueueHookManager get_wrapper failed: {e}") + return result + + execution.PromptQueue.get = get_wrapper + + if self._original_task_done is None: + self._original_task_done = execution.PromptQueue.task_done + + def task_done_wrapper(q_self, item_id, history_result, status): + try: + self._on_task_done((q_self, item_id, history_result, status)) + except Exception as e: + logging.debug(f"QueueHookManager task_done on_task_done failed: {e}") + self._original_task_done(q_self, item_id, history_result, status) + + execution.PromptQueue.task_done = task_done_wrapper + + self._installed = True + + def uninstall(self) -> None: + """Restore original methods.""" + if not self._installed: + return + import execution + if self._original_queue_get is not None: + execution.PromptQueue.get = self._original_queue_get + self._original_queue_get = None + if self._original_task_done is not None: + execution.PromptQueue.task_done = self._original_task_done + self._original_task_done = None + self._installed = False + + diff --git a/server/routes_helper.py b/server/routes_helper.py new file mode 100644 index 0000000..ba0005a --- /dev/null +++ b/server/routes_helper.py @@ -0,0 +1,28 @@ +from typing import Any +from aiohttp import web + + +class RoutesHelper: + """Registers HTTP routes for the persistent queue API.""" + + def __init__(self, app: Any): + self.app = app + + def register(self, manager: "PersistentQueueManager") -> None: + routes = [ + web.get('/api/pqueue', manager._api_get_pqueue), + web.get('/api/pqueue/history', manager._api_get_history), + web.get('/api/pqueue/history/thumb/{history_id:\\d+}', manager._api_get_history_thumb), + web.get('/api/pqueue/preview', manager._api_preview_image), + web.post('/api/pqueue/pause', manager._api_pause), + web.post('/api/pqueue/resume', manager._api_resume), + web.post('/api/pqueue/reorder', manager._api_reorder), + web.patch('/api/pqueue/priority', manager._api_priority), + web.post('/api/pqueue/delete', manager._api_delete), + web.patch('/api/pqueue/rename', manager._api_rename), + web.post('/api/pqueue/run-selected', manager._api_run_selected), + web.post('/api/pqueue/skip-selected', manager._api_skip_selected), + ] + self.app.add_routes(routes) + + diff --git a/server/thumbnail_service.py b/server/thumbnail_service.py new file mode 100644 index 0000000..5684060 --- /dev/null +++ b/server/thumbnail_service.py @@ -0,0 +1,184 @@ +import os +import json +from typing import Optional, Any, Dict, List + +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from PIL import ImageDraw, ImageFont +from io import BytesIO + +import folder_paths + + +class ThumbnailService: + """Generates small, web-friendly thumbnails from ComfyUI output descriptors. + + Single responsibility: image IO and thumbnail encoding. + """ + + def __init__(self, max_size: int = 128, quality: int = 60): + self.max_size = max_size + self.quality = quality + + def generate_thumbnails_from_outputs(self, outputs: Optional[dict], *, workflow_json: Optional[str] = None, extras: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: + if not outputs: + return [] + images: List[Dict[str, Any]] = self._extract_image_descriptors(outputs) + thumbs: List[Dict[str, Any]] = [] + for idx, desc in enumerate(images[:4]): + thumb = self._encode_single_thumbnail(desc, idx, workflow_json=workflow_json, extras=extras) + if thumb is not None: + thumbs.append(thumb) + return thumbs + + def generate_placeholder_thumbnail(self, status: str, *, workflow_json: Optional[str] = None) -> Optional[Dict[str, Any]]: + """Create a placeholder WEBP thumbnail for failed/interrupted jobs. + + The placeholder encodes the workflow in EXIF so it can be restored via drag-and-drop. + """ + try: + size = int(self.max_size) + # Locate provided placeholder asset + placeholder_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'web', 'img', 'failed.png')) + use_draw_fallback = True + if os.path.isfile(placeholder_path): + try: + with Image.open(placeholder_path) as base_img: + # Resize to fit within max size, maintain aspect ratio + w, h = base_img.size + scale = min(size / max(1, w), size / max(1, h), 1.0) + new_size = (max(1, int(w * scale)), max(1, int(h * scale))) + if hasattr(Image, 'Resampling'): + resampling = Image.Resampling.LANCZOS + else: + resampling = Image.LANCZOS + img = base_img.convert('RGB').resize(new_size, resampling) + use_draw_fallback = False + except Exception: + use_draw_fallback = True + + if use_draw_fallback: + # Minimal drawn fallback if asset missing/unreadable + bg = (28, 29, 32) + color = (239, 68, 68) if (status or '').lower() == 'failed' else (234, 179, 8) + text = '×' if (status or '').lower() == 'failed' else '!' + img = Image.new('RGB', (size, size), bg) + draw = ImageDraw.Draw(img) + margin = 12 + rect = (margin, margin, size - margin, size - margin) + try: + draw.rounded_rectangle(rect, radius=10, fill=color) + except Exception: + draw.rectangle(rect, fill=color) + try: + font = ImageFont.load_default() + except Exception: + font = None + if font is not None: + try: + bbox = draw.textbbox((0, 0), text, font=font) + tw = bbox[2] - bbox[0] + th = bbox[3] - bbox[1] + except Exception: + tw, th = draw.textlength(text, font=font), 10 + tx = (size - tw) // 2 + ty = (size - th) // 2 + draw.text((tx, ty), text, font=font, fill=(255, 255, 255)) + + buf = BytesIO() + try: + pnginfo = PngInfo() + if workflow_json: + pnginfo.add_text('prompt', workflow_json) + # Save lossless PNG to preserve colors and avoid artifacts + img.save(buf, format='PNG', pnginfo=pnginfo, compress_level=4) + except Exception: + img.save(buf, format='PNG', compress_level=4) + data = buf.getvalue() + return { + 'idx': 0, + 'mime': 'image/png', + 'width': img.size[0], + 'height': img.size[1], + 'data': data, + } + except Exception: + return None + + def _extract_image_descriptors(self, outputs: dict) -> List[Dict[str, Any]]: + images: List[Dict[str, Any]] = [] + try: + for v in (outputs or {}).values(): + if isinstance(v, dict): + imgs = v.get('images') or (v.get('ui') or {}).get('images') or [] + if isinstance(imgs, list): + for i in imgs: + if isinstance(i, dict) and (i.get('filename') or i.get('name')): + images.append(i) + elif isinstance(v, list): + for i in v: + if isinstance(i, dict) and (i.get('filename') or i.get('name')): + images.append(i) + except Exception: + return [] + return images + + def _encode_single_thumbnail(self, desc: Dict[str, Any], idx: int, *, workflow_json: Optional[str], extras: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + filename = desc.get('filename') or desc.get('name') + folder_type = desc.get('type') or 'output' + subfolder = desc.get('subfolder') or '' + + base_dir = folder_paths.get_directory_by_type(folder_type) + if base_dir is None: + return None + img_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 None + img_dir = full_output_dir + file_path = os.path.join(img_dir, os.path.basename(filename)) + if not os.path.isfile(file_path): + return None + + try: + with Image.open(file_path) as img: + # Resize preserving aspect ratio to fit within max_size + w, h = img.size + scale = min(self.max_size / max(1, w), self.max_size / max(1, h), 1.0) + new_size = (max(1, int(w * scale)), max(1, int(h * scale))) + if hasattr(Image, 'Resampling'): + resampling = Image.Resampling.LANCZOS + else: + resampling = Image.LANCZOS + thumb_img = img.convert('RGB').resize(new_size, resampling) + buf = BytesIO() + # Embed workflow metadata into WEBP EXIF so drag-and-drop works + try: + exif = thumb_img.getexif() + if workflow_json: + # Comfy expects 'prompt:' in 0x0110 for WEBP + exif[0x0110] = "prompt:{}".format(workflow_json) + if extras: + tag = 0x010F + for k, v in extras.items(): + try: + exif[tag] = f"{k}:{json.dumps(v)}" + except Exception: + pass + tag -= 1 + thumb_img.save(buf, format='WEBP', quality=self.quality, exif=exif) + except Exception: + thumb_img.save(buf, format='WEBP', quality=self.quality) + data = buf.getvalue() + return { + 'idx': idx, + 'mime': 'image/webp', + 'width': new_size[0], + 'height': new_size[1], + 'data': data, + } + except Exception: + return None + +