620 lines
26 KiB
Python
620 lines
26 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
import heapq
|
|
import copy
|
|
from typing import Optional, Any, Dict, Tuple, List, Callable
|
|
|
|
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
|
|
import json
|
|
|
|
|
|
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]) -> 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)
|
|
if thumb is not None:
|
|
thumbs.append(thumb)
|
|
return thumbs
|
|
|
|
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) -> 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()
|
|
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]):
|
|
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
|
|
|
|
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 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]
|
|
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),
|
|
]
|
|
self.app.add_routes(routes)
|
|
|
|
class PersistentQueueManager:
|
|
"""Coordinator for persistent queue persistence, API handlers, and hooks."""
|
|
|
|
def __init__(self):
|
|
self.db: QueueDatabase = QueueDatabase()
|
|
self.thumbs: ThumbnailService = ThumbnailService(max_size=128, quality=60)
|
|
self.paused: bool = False
|
|
self.current_job: Optional[Any] = None
|
|
self._installed: bool = False
|
|
self._hooks: Optional[QueueHookManager] = None
|
|
|
|
def initialize(self) -> None:
|
|
"""Install hooks and API routes after PromptServer is created."""
|
|
if self._installed:
|
|
return
|
|
|
|
# Import here to avoid circular dependencies
|
|
from server import PromptServer
|
|
|
|
# Register on-prompt handler to persist incoming prompts
|
|
PromptServer.instance.add_on_prompt_handler(self._on_prompt)
|
|
|
|
# Install queue hooks
|
|
self._hooks = QueueHookManager(
|
|
is_paused_fn=lambda: self.paused,
|
|
on_job_started=lambda prompt_id: self.db.update_job_status(prompt_id, 'running'),
|
|
on_task_done=self._on_task_done_persist,
|
|
)
|
|
self._hooks.install()
|
|
|
|
# Add API routes
|
|
try:
|
|
app = PromptServer.instance.app
|
|
RoutesHelper(app).register(self)
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue add routes failed: {e}")
|
|
|
|
# Restore pending jobs on startup
|
|
self._schedule_restore_pending_jobs()
|
|
|
|
self._installed = True
|
|
|
|
def _on_task_done_persist(self, args: Tuple[Any, Any, Any]) -> None:
|
|
"""Persist history and generate thumbnails before original task_done completes.
|
|
|
|
Args:
|
|
args: Tuple of (q_self, item_id, history_result, status) from task_done wrapper.
|
|
"""
|
|
try:
|
|
q_self, item_id, history_result, status = args
|
|
item = q_self.currently_running.get(item_id)
|
|
if item is None:
|
|
return
|
|
prompt_id = item[1]
|
|
prompt = item[2]
|
|
status_str = status.status_str if status is not None else 'success'
|
|
completed = (status.completed if status is not None else True)
|
|
cancelled = isinstance(status_str, str) and status_str.lower() in ('cancelled', 'canceled', 'interrupted', 'cancel')
|
|
new_state = 'completed'
|
|
if not completed:
|
|
new_state = 'interrupted' if cancelled else 'failed'
|
|
# Persist history and thumbnails BEFORE notifying original logic, to avoid UI race
|
|
self.db.update_job_status(prompt_id, new_state, error=None if completed else status_str)
|
|
history_id = self.db.add_history(
|
|
prompt_id=prompt_id,
|
|
workflow=prompt,
|
|
outputs=history_result.get('outputs', {}),
|
|
status=status_str,
|
|
duration_seconds=None,
|
|
)
|
|
try:
|
|
thumbs = self.thumbs.generate_thumbnails_from_outputs(history_result.get('outputs', {}))
|
|
if thumbs:
|
|
self.db.save_history_thumbnails(history_id, thumbs)
|
|
except Exception as te:
|
|
logging.debug(f"PersistentQueue: failed to save thumbnails: {te}")
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue _on_task_done_persist failed: {e}")
|
|
|
|
def _on_prompt(self, json_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
try:
|
|
if "prompt" in json_data:
|
|
prompt = json_data["prompt"]
|
|
prompt_id = str(json_data.get("prompt_id")) if json_data.get("prompt_id") else None
|
|
if prompt_id is None:
|
|
import uuid
|
|
prompt_id = uuid.uuid4().hex
|
|
json_data["prompt_id"] = prompt_id
|
|
# Priority scaffold (0 default)
|
|
priority = 0
|
|
self.db.add_job(prompt_id, prompt, priority=priority)
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue on_prompt persist failed: {e}")
|
|
return json_data
|
|
|
|
def pause_queue(self):
|
|
"""Pause queue execution"""
|
|
self.paused = True
|
|
|
|
def resume_queue(self):
|
|
"""Resume queue execution"""
|
|
self.paused = False
|
|
|
|
def reorder_job(self, prompt_id: str, new_priority: int):
|
|
"""Change job priority"""
|
|
self.db.update_job_priority(prompt_id, new_priority)
|
|
|
|
def _schedule_restore_pending_jobs(self):
|
|
try:
|
|
from server import PromptServer
|
|
loop = PromptServer.instance.loop
|
|
loop.create_task(self._restore_pending_jobs_async())
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue schedule restore failed: {e}")
|
|
|
|
async def _restore_pending_jobs_async(self):
|
|
from server import PromptServer
|
|
import execution
|
|
server_instance = PromptServer.instance
|
|
|
|
pending_jobs = self.db.get_pending_jobs()
|
|
for job in pending_jobs:
|
|
try:
|
|
prompt_id = job["prompt_id"]
|
|
workflow_json = job["workflow"]
|
|
prompt = json.loads(workflow_json) if isinstance(workflow_json, str) else workflow_json
|
|
valid, err, outputs_to_execute, node_errors = await execution.validate_prompt(prompt_id, prompt, None)
|
|
if valid:
|
|
number = server_instance.number
|
|
server_instance.number += 1
|
|
extra_data = {}
|
|
server_instance.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
|
|
else:
|
|
self.db.update_job_status(prompt_id, 'failed', error=(err or {}).get('message') if isinstance(err, dict) else str(err))
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue restore job failed: {e}")
|
|
|
|
# API Routes
|
|
async def _api_get_pqueue(self, request: web.Request) -> web.Response:
|
|
from server import PromptServer
|
|
running, queued = PromptServer.instance.prompt_queue.get_current_queue_volatile()
|
|
# Compute per-prompt progress (0..1) using global progress registry when available
|
|
progress_map = {}
|
|
try:
|
|
from comfy_execution.progress import get_progress_state
|
|
reg = get_progress_state()
|
|
if reg is not None and getattr(reg, 'nodes', None) and reg.prompt_id:
|
|
total_max = 0.0
|
|
total_val = 0.0
|
|
for st in reg.nodes.values():
|
|
try:
|
|
mv = float(st.get('max', 1.0) or 1.0)
|
|
vv = float(st.get('value', 0.0) or 0.0)
|
|
except Exception:
|
|
mv = 1.0
|
|
vv = 0.0
|
|
if mv <= 0:
|
|
mv = 1.0
|
|
total_max += mv
|
|
total_val += max(0.0, min(vv, mv))
|
|
frac = (total_val / total_max) if total_max > 0 else 0.0
|
|
progress_map[reg.prompt_id] = max(0.0, min(1.0, frac))
|
|
except Exception:
|
|
pass
|
|
|
|
return web.json_response({
|
|
"paused": self.paused,
|
|
"db_pending": self.db.get_pending_jobs(),
|
|
"queue_running": running,
|
|
"queue_pending": queued,
|
|
"running_progress": progress_map,
|
|
})
|
|
|
|
async def _api_get_history(self, request: web.Request) -> web.Response:
|
|
limit = int(request.rel_url.query.get("limit", "50"))
|
|
return web.json_response({"history": self.db.list_history(limit=limit)})
|
|
|
|
# _generate_thumbnails_from_outputs removed in favor of ThumbnailService
|
|
|
|
async def _api_pause(self, request: web.Request) -> web.Response:
|
|
self.pause_queue()
|
|
return web.json_response({"paused": True})
|
|
|
|
async def _api_resume(self, request: web.Request) -> web.Response:
|
|
self.resume_queue()
|
|
return web.json_response({"paused": False})
|
|
|
|
async def _api_reorder(self, request: web.Request) -> web.Response:
|
|
try:
|
|
body = await request.json()
|
|
order: List[str] = body.get("order", [])
|
|
self._rebuild_queue_by_prompt_ids(order)
|
|
return web.json_response({"ok": True})
|
|
except Exception as e:
|
|
logging.warning(f"PersistentQueue reorder failed: {e}")
|
|
return web.json_response({"ok": False, "error": str(e)}, status=400)
|
|
|
|
async def _api_priority(self, request: web.Request) -> web.Response:
|
|
try:
|
|
body = await request.json()
|
|
prompt_id: str = body.get("prompt_id")
|
|
priority: int = int(body.get("priority"))
|
|
if not prompt_id:
|
|
return web.json_response({"ok": False, "error": "prompt_id required"}, status=400)
|
|
self.db.update_job_priority(prompt_id, priority)
|
|
# Apply DB priority to in-memory queue
|
|
self._apply_priority_to_pending()
|
|
return web.json_response({"ok": True})
|
|
except Exception as e:
|
|
logging.warning(f"PersistentQueue set priority failed: {e}")
|
|
return web.json_response({"ok": False, "error": str(e)}, status=400)
|
|
|
|
async def _api_delete(self, request: web.Request) -> web.Response:
|
|
try:
|
|
from server import PromptServer
|
|
body = await request.json()
|
|
prompt_ids: List[str] = body.get("prompt_ids", [])
|
|
q = PromptServer.instance.prompt_queue
|
|
for pid in prompt_ids:
|
|
self.db.remove_job(pid)
|
|
def match(item):
|
|
return item[1] == pid
|
|
q.delete_queue_item(match)
|
|
return web.json_response({"ok": True})
|
|
except Exception as e:
|
|
logging.warning(f"PersistentQueue delete failed: {e}")
|
|
return web.json_response({"ok": False, "error": str(e)}, status=400)
|
|
|
|
def _rebuild_queue_by_prompt_ids(self, ordered_prompt_ids: List[str]) -> None:
|
|
from server import PromptServer
|
|
q = PromptServer.instance.prompt_queue
|
|
with q.mutex:
|
|
# Map current pending items by prompt_id
|
|
by_id: Dict[str, Tuple] = {}
|
|
for item in q.queue:
|
|
# item: (number, prompt_id, prompt, extra_data, outputs_to_execute)
|
|
by_id[item[1]] = item
|
|
|
|
# Build new ordered list; preserve items not specified at the end by current order
|
|
specified = [pid for pid in ordered_prompt_ids if pid in by_id]
|
|
unspecified = [it for it in q.queue if it[1] not in set(specified)]
|
|
|
|
# Assign new numbers so the new order executes first
|
|
new_items: List[Tuple] = []
|
|
next_num = -len(specified) if len(specified) > 0 else 0
|
|
for pid in specified:
|
|
old = by_id[pid]
|
|
new_items.append((next_num, old[1], old[2], old[3], old[4]))
|
|
next_num += 1
|
|
|
|
# Append unspecified with their existing relative order but after specified
|
|
base = next_num
|
|
for idx, old in enumerate(sorted(unspecified, key=lambda x: x[0])):
|
|
new_items.append((base + idx, old[1], old[2], old[3], old[4]))
|
|
|
|
q.queue = new_items
|
|
heapq.heapify(q.queue)
|
|
q.server.queue_updated()
|
|
|
|
def _apply_priority_to_pending(self) -> None:
|
|
"""Rebuild in-memory queue using DB priority DESC, then created_at ASC."""
|
|
from server import PromptServer
|
|
q = PromptServer.instance.prompt_queue
|
|
pending = self.db.get_pending_jobs()
|
|
# order prompt_ids by priority desc, created_at asc
|
|
ordered_ids = [row["prompt_id"] for row in sorted(pending, key=lambda r: (-int(r.get("priority", 0)), r.get("created_at") or ""))]
|
|
self._rebuild_queue_by_prompt_ids(ordered_ids)
|
|
|
|
def _install_api_routes(self):
|
|
try:
|
|
from server import PromptServer
|
|
app = PromptServer.instance.app
|
|
app.add_routes([
|
|
web.get('/api/pqueue', self._api_get_pqueue),
|
|
web.get('/api/pqueue/history', self._api_get_history),
|
|
web.get('/api/pqueue/history/thumb/{history_id:\\d+}', self._api_get_history_thumb),
|
|
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),
|
|
web.patch('/api/pqueue/priority', self._api_priority),
|
|
web.post('/api/pqueue/delete', self._api_delete),
|
|
])
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue add routes failed: {e}")
|
|
|
|
async def _api_get_history_thumb(self, request: web.Request) -> web.Response:
|
|
try:
|
|
history_id = int(request.match_info.get('history_id', '0'))
|
|
idx = int(request.rel_url.query.get('idx', '0'))
|
|
row = self.db.get_history_thumbnail(history_id, idx)
|
|
if not row:
|
|
return web.Response(status=404)
|
|
return web.Response(body=row['data'], content_type=row.get('mime') or 'image/webp')
|
|
except Exception:
|
|
return web.Response(status=500)
|
|
|
|
async def _api_preview_image(self, request: web.Request) -> web.Response:
|
|
"""Serve cached previews with embedded workflow metadata if available.
|
|
|
|
Splits concerns into small helpers to keep logic simple (KISS/DRY).
|
|
"""
|
|
try:
|
|
params = self._parse_preview_params(request)
|
|
file = self._resolve_preview_filepath(params)
|
|
if file is None:
|
|
return web.Response(status=400)
|
|
if not os.path.isfile(file):
|
|
return web.Response(status=404)
|
|
|
|
cache_path = self._compute_preview_cache_path(file, params)
|
|
if cache_path and os.path.exists(cache_path):
|
|
return web.FileResponse(cache_path, headers={"Content-Disposition": f"filename=\"{params['filename']}\""})
|
|
|
|
# Render and cache
|
|
cache_path = self._render_and_cache_preview(file, params, cache_path)
|
|
return web.FileResponse(cache_path, headers={"Content-Disposition": f"filename=\"{params['filename']}\""})
|
|
except Exception:
|
|
return web.Response(status=500)
|
|
|
|
def _parse_preview_params(self, request: web.Request) -> Dict[str, Any]:
|
|
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])
|
|
|
|
pid = request.rel_url.query.get('pid')
|
|
return {
|
|
'filename': request.rel_url.query.get('filename'),
|
|
'subfolder': request.rel_url.query.get('subfolder', ''),
|
|
'type': request.rel_url.query.get('type', 'output'),
|
|
'pid': pid,
|
|
'image_format': image_format,
|
|
'quality': quality,
|
|
'workflow_json': self._lookup_workflow_json(pid) if pid else None,
|
|
}
|
|
|
|
def _lookup_workflow_json(self, pid: Optional[str]) -> Optional[str]:
|
|
if not pid:
|
|
return None
|
|
try:
|
|
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']:
|
|
return row['workflow']
|
|
except Exception:
|
|
pass
|
|
job = self.db.get_job(pid)
|
|
if job:
|
|
return job.get('workflow')
|
|
return None
|
|
|
|
def _resolve_preview_filepath(self, params: Dict[str, Any]) -> Optional[str]:
|
|
base_dir = folder_paths.get_directory_by_type(params.get('type') or 'output')
|
|
if base_dir is None:
|
|
return None
|
|
output_dir = base_dir
|
|
subfolder = params.get('subfolder') or ''
|
|
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
|
|
output_dir = full_output_dir
|
|
filename = os.path.basename(params.get('filename') or '')
|
|
return os.path.join(output_dir, filename)
|
|
|
|
def _compute_preview_cache_path(self, file: str, params: Dict[str, Any]) -> Optional[str]:
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
cache_dir = os.path.join(temp_dir, 'preview_cache')
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
wf_hash = hashlib.sha256(((params.get('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}|{params['image_format']}|{params['quality']}|{wf_hash}"
|
|
cache_name = hashlib.sha256(cache_key.encode('utf-8')).hexdigest() + f'.{params["image_format"]}'
|
|
return os.path.join(cache_dir, cache_name)
|
|
|
|
def _render_and_cache_preview(self, file: str, params: Dict[str, Any], cache_path: Optional[str]) -> str:
|
|
with Image.open(file) as img:
|
|
save_kwargs: Dict[str, Any] = {"format": params['image_format']}
|
|
if params['image_format'] in ['webp', 'jpeg']:
|
|
save_kwargs['quality'] = params['quality']
|
|
|
|
# Embed metadata
|
|
workflow_json = params.get('workflow_json')
|
|
if params['image_format'] == 'webp':
|
|
self._embed_webp_metadata(img, save_kwargs, workflow_json)
|
|
elif params['image_format'] == 'png':
|
|
self._embed_png_metadata(img, save_kwargs, workflow_json)
|
|
|
|
img.save(cache_path, **save_kwargs)
|
|
return cache_path
|
|
|
|
def _embed_webp_metadata(self, img: Image.Image, save_kwargs: Dict[str, Any], workflow_json: Optional[str]) -> None:
|
|
try:
|
|
exif = img.getexif()
|
|
if hasattr(img, 'text'):
|
|
for k in img.text:
|
|
val = img.text[k]
|
|
if k == 'prompt':
|
|
exif[0x0110] = "prompt:{}".format(val)
|
|
else:
|
|
tag = 0x010F
|
|
try:
|
|
exif[tag] = f"{k}:{val}"
|
|
except Exception:
|
|
pass
|
|
if workflow_json:
|
|
try:
|
|
exif[0x010E] = "workflow:{}".format(workflow_json)
|
|
except Exception:
|
|
pass
|
|
save_kwargs['exif'] = exif
|
|
except Exception:
|
|
pass
|
|
|
|
def _embed_png_metadata(self, img: Image.Image, save_kwargs: Dict[str, Any], workflow_json: Optional[str]) -> None:
|
|
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
|
|
|
|
queue_manager = PersistentQueueManager() |