Refactor: Split server module into managable chunks
This commit is contained in:
Simon Cordt
+3
-2
@@ -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}")
|
||||
|
||||
|
||||
@@ -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:<json>' 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()
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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:<json>' 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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user