diff --git a/server/queue_manager.py b/server/queue_manager.py index affc87b..e2bd1a1 100644 --- a/server/queue_manager.py +++ b/server/queue_manager.py @@ -16,6 +16,7 @@ import os import hashlib import folder_paths import json +from PIL import ImageDraw, ImageFont class ThumbnailService: @@ -28,17 +29,91 @@ class ThumbnailService: self.max_size = max_size self.quality = quality - def generate_thumbnails_from_outputs(self, outputs: Optional[dict]) -> List[Dict[str, Any]]: + 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) + 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: @@ -57,7 +132,7 @@ class ThumbnailService: return [] return images - def _encode_single_thumbnail(self, desc: Dict[str, Any], idx: int) -> Optional[Dict[str, Any]]: + 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 '' @@ -87,7 +162,23 @@ class ThumbnailService: resampling = Image.LANCZOS thumb_img = img.convert('RGB').resize(new_size, resampling) buf = BytesIO() - thumb_img.save(buf, format='WEBP', quality=self.quality) + # 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, @@ -256,7 +347,13 @@ class PersistentQueueManager: duration_seconds=None, ) try: - thumbs = self.thumbs.generate_thumbnails_from_outputs(history_result.get('outputs', {})) + workflow_json = json.dumps(prompt) if isinstance(prompt, (dict, list)) else str(prompt) + outputs = history_result.get('outputs', {}) + thumbs = self.thumbs.generate_thumbnails_from_outputs(outputs, workflow_json=workflow_json, extras=None) + if not thumbs: + ph = self.thumbs.generate_placeholder_thumbnail(new_state, workflow_json=workflow_json) + if ph: + thumbs = [ph] if thumbs: self.db.save_history_thumbnails(history_id, thumbs) except Exception as te: