1098 lines
49 KiB
Python
1098 lines
49 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
import heapq
|
|
from typing import Optional, Any, Dict, Tuple, List, Callable, Set
|
|
|
|
from aiohttp import web
|
|
|
|
from .database import QueueDatabase
|
|
from PIL import Image
|
|
import os
|
|
import hashlib
|
|
import folder_paths
|
|
|
|
from .thumbnail_service import ThumbnailService
|
|
from .queue_hook_manager import QueueHookManager
|
|
from .routes_helper import RoutesHelper
|
|
|
|
|
|
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)
|
|
# Default to paused state on startup for safety - user can resume when ready
|
|
self.paused: bool = True
|
|
self.current_job: Optional[Any] = None
|
|
self._installed: bool = False
|
|
self._hooks: Optional[QueueHookManager] = None
|
|
# IDs allowed to run while paused (run-selected mode)
|
|
self._run_selected_remaining: Set[str] = set()
|
|
# Accumulate normalized progress across multiple sampler nodes per prompt
|
|
self._progress_accum: Dict[str, float] = {}
|
|
self._progress_last_raw: Dict[str, float] = {}
|
|
self._samplers_total: Dict[str, int] = {}
|
|
|
|
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._on_job_started(prompt_id),
|
|
on_task_done=self._on_task_done_persist,
|
|
should_run_when_paused=self._is_prompt_allowed_while_paused,
|
|
)
|
|
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
|
|
|
|
# Log initial state
|
|
logging.info("PersistentQueue initialized in PAUSED state. Use UI to resume queue processing.")
|
|
|
|
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)
|
|
# Ensure any user-provided rename is reflected in the workflow stored to history
|
|
try:
|
|
if isinstance(prompt, dict):
|
|
# If DB has a more recent renamed workflow JSON, prefer its name field(s)
|
|
db_row = self.db.get_job(prompt_id)
|
|
if db_row and db_row.get('workflow'):
|
|
try:
|
|
db_wf = json.loads(db_row['workflow']) if isinstance(db_row['workflow'], str) else db_row['workflow']
|
|
except Exception:
|
|
db_wf = None
|
|
if isinstance(db_wf, dict):
|
|
def _extract_name(wf):
|
|
try:
|
|
return (wf.get('workflow') or {}).get('name') or wf.get('name')
|
|
except Exception:
|
|
return None
|
|
db_name = _extract_name(db_wf)
|
|
if isinstance(db_name, str) and db_name.strip():
|
|
if isinstance(prompt.get('workflow'), dict):
|
|
prompt['workflow']['name'] = db_name.strip()
|
|
else:
|
|
prompt['name'] = db_name.strip()
|
|
except Exception:
|
|
pass
|
|
|
|
history_id = self.db.add_history(
|
|
prompt_id=prompt_id,
|
|
workflow=prompt,
|
|
outputs=history_result.get('outputs', {}),
|
|
status=status_str,
|
|
duration_seconds=None,
|
|
)
|
|
try:
|
|
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:
|
|
logging.debug(f"PersistentQueue: failed to save thumbnails: {te}")
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue _on_task_done_persist failed: {e}")
|
|
# After persisting, if we are in run-selected mode, update remaining set
|
|
try:
|
|
q_self, item_id, history_result, status = args
|
|
item = q_self.currently_running.get(item_id)
|
|
if item is not None:
|
|
pid = str(item[1])
|
|
if pid in self._run_selected_remaining:
|
|
self._run_selected_remaining.discard(pid)
|
|
if not self._run_selected_remaining:
|
|
# All selected jobs finished; keep queue paused and clear state
|
|
logging.info("PersistentQueue: Finished all selected jobs; queue remains paused.")
|
|
except Exception:
|
|
pass
|
|
|
|
def _on_job_started(self, prompt_id: str):
|
|
"""Called when a job transitions to running status."""
|
|
try:
|
|
# Update job status in database
|
|
self.db.update_job_status(prompt_id, 'running')
|
|
|
|
# Clear cached sampler count to force recalculation with fresh data
|
|
if hasattr(self, '_samplers_total') and isinstance(self._samplers_total, dict):
|
|
self._samplers_total.pop(str(prompt_id), None)
|
|
|
|
except Exception as e:
|
|
logging.debug(f"PersistentQueue _on_job_started 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
|
|
# Reset any cached progress/sampler state for this prompt id (new run)
|
|
try:
|
|
if hasattr(self, '_progress_accum') and isinstance(self._progress_accum, dict):
|
|
self._progress_accum.pop(str(prompt_id), None)
|
|
if hasattr(self, '_progress_last_raw') and isinstance(self._progress_last_raw, dict):
|
|
self._progress_last_raw.pop(str(prompt_id), None)
|
|
if hasattr(self, '_samplers_total') and isinstance(self._samplers_total, dict):
|
|
self._samplers_total.pop(str(prompt_id), None)
|
|
except Exception:
|
|
pass
|
|
# Optional: extract suggested name from extra_data
|
|
try:
|
|
extra = json_data.get('extra_data') or {}
|
|
name_hint = None
|
|
if isinstance(extra, dict):
|
|
name_hint = extra.get('pqueue_workflow_name')
|
|
# Build a persistence-only copy with the name; DO NOT mutate json_data['prompt']
|
|
persist_prompt = prompt
|
|
if isinstance(name_hint, str) and name_hint.strip() and isinstance(prompt, dict):
|
|
if isinstance(prompt.get('workflow'), dict):
|
|
persist_prompt = dict(prompt)
|
|
wf = dict(persist_prompt['workflow'])
|
|
wf['name'] = name_hint.strip()
|
|
persist_prompt['workflow'] = wf
|
|
else:
|
|
persist_prompt = dict(prompt)
|
|
persist_prompt['name'] = name_hint.strip()
|
|
else:
|
|
persist_prompt = prompt
|
|
except Exception:
|
|
persist_prompt = prompt
|
|
|
|
# Priority scaffold (0 default)
|
|
priority = 0
|
|
self.db.add_job(prompt_id, persist_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()
|
|
logging.info(f"PersistentQueue: Found {len(pending_jobs)} pending jobs to restore on startup")
|
|
restored_count = 0
|
|
failed_restores = []
|
|
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
|
|
# Robust normalization: clean extraneous keys and extract inner workflow if present
|
|
if isinstance(prompt, dict):
|
|
prompt.pop('name', None) # Remove top-level name (non-node key)
|
|
if 'workflow' in prompt and isinstance(prompt['workflow'], dict):
|
|
inner = prompt['workflow']
|
|
inner.pop('name', None) # Clean name from inner if added
|
|
prompt = inner
|
|
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))
|
|
restored_count += 1
|
|
logging.info(f"PersistentQueue: Successfully restored job {prompt_id} to queue")
|
|
else:
|
|
error_msg = (err or {}).get('message') if isinstance(err, dict) else str(err)
|
|
self.db.update_job_status(prompt_id, 'failed', error=error_msg)
|
|
failed_restores.append((prompt_id, error_msg))
|
|
logging.warning(f"PersistentQueue: Failed to restore job {prompt_id}: {error_msg}")
|
|
except Exception as e:
|
|
logging.error(f"PersistentQueue: Exception during restore of job {job.get('prompt_id', 'unknown')}: {e}")
|
|
failed_restores.append((job.get('prompt_id', 'unknown'), str(e)))
|
|
|
|
if restored_count > 0:
|
|
logging.info(f"PersistentQueue: Restored {restored_count} jobs to queue")
|
|
if failed_restores:
|
|
logging.warning(f"PersistentQueue: {len(failed_restores)} jobs failed to restore: {failed_restores}")
|
|
if len(pending_jobs) == 0:
|
|
logging.info("PersistentQueue: No pending jobs to restore on startup")
|
|
|
|
# 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()
|
|
# Ensure queued list is sorted by execution order (heap array is not fully ordered)
|
|
try:
|
|
queued_sorted = sorted(queued, key=lambda it: (it[0], str(it[1])))
|
|
except Exception:
|
|
queued_sorted = queued
|
|
# Compute per-prompt progress on client via sockets; keep server map minimal
|
|
progress_map = {}
|
|
# Build quick lookup for running prompts
|
|
running_prompts: Dict[str, Any] = {}
|
|
try:
|
|
for it in running or []:
|
|
try:
|
|
pid = str(it[1])
|
|
running_prompts[pid] = it[2]
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
# Clear accumulators for prompts that are no longer running
|
|
try:
|
|
current_running_ids = set(running_prompts.keys())
|
|
for pid in list(self._progress_accum.keys()):
|
|
if pid not in current_running_ids:
|
|
try:
|
|
del self._progress_accum[pid]
|
|
except Exception:
|
|
pass
|
|
try:
|
|
del self._progress_last_raw[pid]
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
# Intentionally avoid server-side progress normalization; client will handle via sockets
|
|
|
|
# Provide sampler counts for running prompts so frontend can normalize socket progress
|
|
sampler_count_by_id: Dict[str, int] = {}
|
|
try:
|
|
for pid, prompt in running_prompts.items():
|
|
try:
|
|
# Always fetch fresh workflow from DB for running jobs to avoid stale prompt data
|
|
# This ensures we get the correct sampler count even if the workflow was modified
|
|
prompt = running_prompts[pid]
|
|
sampler_count_by_id[pid] = int(self._get_total_samplers(pid, prompt))
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
return web.json_response({
|
|
"paused": self.paused,
|
|
"db_pending": self.db.get_pending_jobs(),
|
|
"queue_running": running,
|
|
"queue_pending": queued_sorted,
|
|
"running_progress": progress_map,
|
|
"sampler_count_by_id": sampler_count_by_id,
|
|
# Provide DB rows for ALL visible queue items (pending + running) so UI
|
|
# can derive labels (including renamed names) even after status changes
|
|
"db_by_id": self._build_db_lookup_for_queue_items(running, queued_sorted),
|
|
})
|
|
|
|
async def _api_export_queue(self, request: web.Request) -> web.StreamResponse:
|
|
"""Export the current queue (pending items) as ordered JSON.
|
|
|
|
Format:
|
|
{
|
|
"version": 1,
|
|
"exported_at": <iso string>,
|
|
"items": [
|
|
{ "prompt_id": str, "workflow": object, "priority": int, "created_at": ts/string }
|
|
]
|
|
}
|
|
"""
|
|
try:
|
|
from server import PromptServer
|
|
import datetime
|
|
# Build ordered pending list by the actual execution order in in-memory queue
|
|
q = PromptServer.instance.prompt_queue
|
|
with q.mutex:
|
|
try:
|
|
queued = list(q.queue)
|
|
except Exception:
|
|
queued = []
|
|
# Sort by number to reflect real order (lowest number executes first)
|
|
try:
|
|
queued_sorted = sorted(queued, key=lambda it: (it[0], str(it[1])))
|
|
except Exception:
|
|
queued_sorted = queued
|
|
|
|
# Build DB lookup to pull stored workflow/priority/created_at
|
|
db_lookup = self._build_db_lookup_for_queue_items([], queued_sorted)
|
|
|
|
items = []
|
|
for it in queued_sorted:
|
|
try:
|
|
pid = str(it[1])
|
|
db_row = db_lookup.get(pid) or {}
|
|
wf_text = db_row.get('workflow')
|
|
try:
|
|
workflow = json.loads(wf_text) if isinstance(wf_text, str) else (wf_text if wf_text is not None else None)
|
|
except Exception:
|
|
workflow = None
|
|
items.append({
|
|
"prompt_id": pid,
|
|
"workflow": workflow,
|
|
"priority": int(db_row.get('priority') or 0),
|
|
"created_at": db_row.get('created_at'),
|
|
})
|
|
except Exception:
|
|
pass
|
|
|
|
payload = {
|
|
"version": 1,
|
|
"exported_at": datetime.datetime.now().isoformat(sep=' ', timespec='seconds'),
|
|
"items": items,
|
|
}
|
|
|
|
text = json.dumps(payload, ensure_ascii=False)
|
|
return web.Response(
|
|
text=text,
|
|
content_type='application/json',
|
|
headers={
|
|
'Content-Disposition': 'attachment; filename="pqueue-export.json"'
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logging.warning(f"PersistentQueue export failed: {e}")
|
|
return web.json_response({"ok": False, "error": str(e)}, status=500)
|
|
|
|
def _build_db_lookup_for_queue_items(self, running: List[Tuple], queued: List[Tuple]) -> Dict[str, Dict[str, Any]]:
|
|
"""Return a mapping of prompt_id -> DB row for all running and queued items.
|
|
|
|
Includes at least the workflow text so the UI can derive renamed labels.
|
|
"""
|
|
try:
|
|
pids_set: Set[str] = set()
|
|
try:
|
|
for it in running or []:
|
|
try:
|
|
pids_set.add(str(it[1]))
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
try:
|
|
for it in queued or []:
|
|
try:
|
|
pids_set.add(str(it[1]))
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
pids = [pid for pid in pids_set if pid]
|
|
if not pids:
|
|
return {}
|
|
rows: Dict[str, Dict[str, Any]] = {}
|
|
# Prefer batched queries to avoid SQLite var limits; fall back to per-id on error
|
|
try:
|
|
with self.db._get_conn() as conn:
|
|
batch_size = 500
|
|
for i in range(0, len(pids), batch_size):
|
|
batch = pids[i:i+batch_size]
|
|
placeholders = ",".join(["?"] * len(batch))
|
|
try:
|
|
cur = conn.execute(f"SELECT * FROM queue_items WHERE prompt_id IN ({placeholders})", tuple(batch))
|
|
for r in cur.fetchall():
|
|
try:
|
|
d = dict(r)
|
|
pid = str(d.get('prompt_id')) if d.get('prompt_id') is not None else None
|
|
if pid:
|
|
rows[pid] = d
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
# Fallback to per-id lookup for this batch if SQL fails
|
|
for pid in batch:
|
|
try:
|
|
row = self.db.get_job(pid)
|
|
if row:
|
|
rows[pid] = row
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
# Fallback to per-id for all if connection or other errors
|
|
for pid in pids:
|
|
try:
|
|
row = self.db.get_job(pid)
|
|
if row:
|
|
rows[pid] = row
|
|
except Exception:
|
|
pass
|
|
return rows
|
|
except Exception:
|
|
return {}
|
|
|
|
async def _api_get_history(self, request: web.Request) -> web.Response:
|
|
# Support both simple and paginated calls. Frontend can pass sort/filter + cursor for keyset pagination.
|
|
q = request.rel_url.query
|
|
try:
|
|
limit = int(q.get("limit", "50"))
|
|
except Exception:
|
|
limit = 50
|
|
|
|
sort_by = q.get("sort_by") or "id"
|
|
sort_dir = q.get("sort_dir") or "desc"
|
|
|
|
# Cursor-based pagination
|
|
cursor_id = None
|
|
cursor_value = None
|
|
try:
|
|
if q.get("cursor_id") is not None:
|
|
cursor_id = int(q.get("cursor_id"))
|
|
except Exception:
|
|
cursor_id = None
|
|
if q.get("cursor_value") is not None:
|
|
cursor_value = q.get("cursor_value")
|
|
|
|
status = q.get("status")
|
|
search_q = q.get("q")
|
|
since = q.get("since")
|
|
until = q.get("until")
|
|
try:
|
|
min_duration = float(q.get("min_duration")) if q.get("min_duration") is not None else None
|
|
except Exception:
|
|
min_duration = None
|
|
try:
|
|
max_duration = float(q.get("max_duration")) if q.get("max_duration") is not None else None
|
|
except Exception:
|
|
max_duration = None
|
|
|
|
# If only limit is provided and no advanced params, keep legacy behavior
|
|
legacy_mode = (
|
|
(q.keys() <= {"limit"}) or
|
|
(set(q.keys()) == set() )
|
|
)
|
|
if legacy_mode:
|
|
return web.json_response({"history": self.db.list_history(limit=limit)})
|
|
|
|
result = self.db.list_history_paginated(
|
|
limit=limit,
|
|
sort_by=sort_by,
|
|
sort_dir=sort_dir,
|
|
cursor_id=cursor_id,
|
|
cursor_value=cursor_value,
|
|
status=status,
|
|
q=search_q,
|
|
since=since,
|
|
until=until,
|
|
min_duration=min_duration,
|
|
max_duration=max_duration,
|
|
)
|
|
return web.json_response(result)
|
|
|
|
# _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_import_queue(self, request: web.Request) -> web.Response:
|
|
"""Import a queue JSON file and append items after existing ones, preserving order.
|
|
|
|
Accepts either application/json body or multipart/form-data with a single file field named 'file'.
|
|
"""
|
|
try:
|
|
# Parse JSON payload (supports both raw JSON and multipart file upload containing JSON)
|
|
data = None
|
|
ctype = (request.headers.get('Content-Type') or '').lower()
|
|
if 'application/json' in ctype:
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
data = None
|
|
if data is None:
|
|
# Try multipart
|
|
reader = await request.multipart()
|
|
if reader is not None:
|
|
async for part in reader:
|
|
if part.name == 'file':
|
|
try:
|
|
raw = await part.read(decode=True)
|
|
except Exception:
|
|
raw = await part.read()
|
|
try:
|
|
text = raw.decode('utf-8', errors='replace') if isinstance(raw, (bytes, bytearray)) else str(raw)
|
|
data = json.loads(text)
|
|
except Exception:
|
|
data = None
|
|
break
|
|
|
|
if not isinstance(data, dict):
|
|
return web.json_response({"ok": False, "error": "Invalid JSON"}, status=400)
|
|
items = data.get('items')
|
|
if not isinstance(items, list):
|
|
return web.json_response({"ok": False, "error": "Missing items[]"}, status=400)
|
|
|
|
# Normalize imported items and enqueue after existing ones
|
|
from server import PromptServer
|
|
import execution
|
|
q = PromptServer.instance.prompt_queue
|
|
|
|
# Snapshot current queue length and max number to compute appended numbers
|
|
with q.mutex:
|
|
try:
|
|
existing = list(q.queue)
|
|
except Exception:
|
|
existing = []
|
|
try:
|
|
max_num = max((it[0] for it in existing), default=0)
|
|
except Exception:
|
|
max_num = 0
|
|
|
|
appended = []
|
|
server_instance = PromptServer.instance
|
|
for item in items:
|
|
try:
|
|
pid = str(item.get('prompt_id')) if item.get('prompt_id') is not None else None
|
|
workflow = item.get('workflow')
|
|
if not pid or workflow is None:
|
|
continue
|
|
# De-duplicate: skip if already present in in-memory queue
|
|
present = False
|
|
with q.mutex:
|
|
try:
|
|
present = any(str(it[1]) == pid for it in q.queue)
|
|
except Exception:
|
|
present = False
|
|
if present:
|
|
continue
|
|
|
|
# Persist to DB (if not already)
|
|
try:
|
|
self.db.add_job(pid, workflow, priority=int(item.get('priority') or 0))
|
|
except Exception:
|
|
pass
|
|
|
|
# Validate normalized prompt for execution
|
|
try:
|
|
prompt = workflow
|
|
if isinstance(prompt, dict):
|
|
# Clean metadata fields that are not part of executable nodes
|
|
prompt = dict(prompt)
|
|
prompt.pop('name', None)
|
|
if 'workflow' in prompt and isinstance(prompt['workflow'], dict):
|
|
inner = dict(prompt['workflow'])
|
|
inner.pop('name', None)
|
|
prompt = inner
|
|
except Exception:
|
|
prompt = workflow
|
|
|
|
valid, err, outputs_to_execute, node_errors = await execution.validate_prompt(pid, prompt, None)
|
|
if not valid:
|
|
# Mark failed import for this id but continue others
|
|
try:
|
|
self.db.update_job_status(pid, 'failed', error=(err or {}).get('message') if isinstance(err, dict) else str(err))
|
|
except Exception:
|
|
pass
|
|
continue
|
|
|
|
# Assign increasing numbers after current max to append at end
|
|
with q.mutex:
|
|
number = max_num + 1
|
|
max_num = number
|
|
extra_data = {}
|
|
q.queue.append((number, pid, prompt, extra_data, outputs_to_execute))
|
|
heapq.heapify(q.queue)
|
|
try:
|
|
q.server.queue_updated()
|
|
except Exception:
|
|
pass
|
|
appended.append(pid)
|
|
except Exception:
|
|
continue
|
|
|
|
return web.json_response({"ok": True, "imported": appended, "count": len(appended)})
|
|
except Exception as e:
|
|
logging.warning(f"PersistentQueue import failed: {e}")
|
|
return web.json_response({"ok": False, "error": str(e)}, status=500)
|
|
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)
|
|
|
|
async def _api_rename(self, request: web.Request) -> web.Response:
|
|
try:
|
|
from server import PromptServer
|
|
body = await request.json()
|
|
prompt_id: str = body.get('prompt_id')
|
|
new_name: str = body.get('name')
|
|
if not prompt_id or new_name is None:
|
|
return web.json_response({"ok": False, "error": "prompt_id and name required"}, status=400)
|
|
# Persist to DB
|
|
ok = self.db.update_job_name(prompt_id, str(new_name))
|
|
if not ok:
|
|
return web.json_response({"ok": False, "error": "job not found"}, status=404)
|
|
# Do NOT mutate in-memory prompt JSON in the queue. The UI derives names from DB.
|
|
# We intentionally avoid adding non-node keys (e.g. name/workflow) to the prompt to prevent execution errors.
|
|
return web.json_response({"ok": True})
|
|
except Exception as e:
|
|
logging.warning(f"PersistentQueue rename failed: {e}")
|
|
return web.json_response({"ok": False, "error": str(e)}, status=400)
|
|
|
|
async def _api_run_selected(self, request: web.Request) -> web.Response:
|
|
"""Run selected jobs when queue is paused"""
|
|
try:
|
|
from server import PromptServer
|
|
import execution
|
|
body = await request.json()
|
|
prompt_ids: List[str] = body.get('prompt_ids', [])
|
|
|
|
logging.info(f"PersistentQueue: run_selected called with {len(prompt_ids)} jobs: {prompt_ids}")
|
|
logging.info(f"PersistentQueue: paused state = {self.paused}")
|
|
|
|
if not self.paused:
|
|
return web.json_response({"ok": False, "error": "Queue must be paused to run selected jobs"}, status=400)
|
|
|
|
q = PromptServer.instance.prompt_queue
|
|
server_instance = PromptServer.instance
|
|
|
|
# First, ensure selected jobs are properly loaded into the queue
|
|
executed_ids = []
|
|
|
|
# Prepare selected items and determine order by queue number (ascending = top-most first)
|
|
missing_ids: List[str] = []
|
|
with q.mutex:
|
|
# Map current queue items by prompt_id
|
|
queue_by_id = {str(item[1]): item for item in q.queue}
|
|
|
|
selected_set = set(map(str, prompt_ids))
|
|
# Sort selected that are present by their current number descending (newest/bottom first)
|
|
selected_present = [queue_by_id[pid] for pid in selected_set if pid in queue_by_id]
|
|
selected_present_sorted = sorted(selected_present, key=lambda it: it[0])
|
|
executed_ids = [str(it[1]) for it in selected_present_sorted]
|
|
|
|
# Any selected IDs not in queue are missing; handle after lock
|
|
missing_ids = [pid for pid in map(str, prompt_ids) if pid not in queue_by_id]
|
|
|
|
# Load and validate any missing selected from DB (outside of queue lock)
|
|
for prompt_id in missing_ids:
|
|
job = self.db.get_job(prompt_id)
|
|
if not job:
|
|
continue
|
|
try:
|
|
workflow_json = job["workflow"]
|
|
prompt = json.loads(workflow_json) if isinstance(workflow_json, str) else workflow_json
|
|
# Clean rename metadata
|
|
if isinstance(prompt, dict):
|
|
if 'workflow' in prompt and isinstance(prompt['workflow'], dict):
|
|
prompt.pop('workflow', None)
|
|
if 'name' in prompt and isinstance(prompt['name'], str):
|
|
prompt.pop('name', None)
|
|
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 = {}
|
|
item = (number, prompt_id, prompt, extra_data, outputs_to_execute)
|
|
# Insert this new item at end for now; reorder shortly
|
|
with q.mutex:
|
|
q.queue.append(item)
|
|
heapq.heapify(q.queue)
|
|
executed_ids.append(prompt_id)
|
|
else:
|
|
logging.warning(f"PersistentQueue: Invalid prompt {prompt_id}: {err}")
|
|
except Exception as e:
|
|
logging.error(f"PersistentQueue: Failed to load job {prompt_id}: {e}")
|
|
|
|
# Rebuild queue with selected items first, preserving their relative order by number (desc)
|
|
if executed_ids:
|
|
with q.mutex:
|
|
# Build final selected from current queue and sort by number asc to preserve queue order
|
|
current_by_id = {str(item[1]): item for item in q.queue}
|
|
selected_final = [current_by_id[pid] for pid in executed_ids if pid in current_by_id]
|
|
selected_final_sorted = sorted(selected_final, key=lambda it: it[0])
|
|
selected_ids_final = [str(it[1]) for it in selected_final_sorted]
|
|
self._rebuild_queue_by_prompt_ids(selected_ids_final)
|
|
|
|
# Enable run-selected mode while keeping queue paused so only selected items run
|
|
self._run_selected_remaining = set(map(str, executed_ids))
|
|
logging.info(f"PersistentQueue: Run-selected mode enabled for {len(executed_ids)} jobs; queue remains paused")
|
|
|
|
return web.json_response({"ok": True, "executed": executed_ids})
|
|
except Exception as e:
|
|
logging.error(f"PersistentQueue run selected failed: {e}", exc_info=True)
|
|
return web.json_response({"ok": False, "error": str(e)}, status=500)
|
|
|
|
async def _api_skip_selected(self, request: web.Request) -> web.Response:
|
|
"""Skip selected jobs (remove from queue)"""
|
|
try:
|
|
from server import PromptServer
|
|
body = await request.json()
|
|
prompt_ids: List[str] = body.get('prompt_ids', [])
|
|
|
|
if self.paused:
|
|
return web.json_response({"ok": False, "error": "Queue must be running to skip jobs"}, status=400)
|
|
|
|
q = PromptServer.instance.prompt_queue
|
|
skipped_ids = []
|
|
|
|
for pid in prompt_ids:
|
|
try:
|
|
q.delete(pid)
|
|
self.db.update_job_status(pid, 'cancelled')
|
|
skipped_ids.append(pid)
|
|
except Exception:
|
|
pass
|
|
|
|
return web.json_response({"ok": True, "skipped": skipped_ids})
|
|
except Exception as e:
|
|
logging.warning(f"PersistentQueue skip selected failed: {e}")
|
|
return web.json_response({"ok": False, "error": str(e)}, status=500)
|
|
|
|
def _rebuild_queue_by_prompt_ids(self, selected_ids_in_order: List[str]) -> None:
|
|
"""Promote only the selected prompt_ids to the front and leave others unchanged.
|
|
|
|
The selected_ids_in_order must be in desired execution order (first executes first).
|
|
We assign them numbers lower than any existing number to ensure they run next,
|
|
while keeping all other items' numbers intact to preserve global order.
|
|
"""
|
|
from server import PromptServer
|
|
q = PromptServer.instance.prompt_queue
|
|
with q.mutex:
|
|
# Normalize to strings for safe comparisons
|
|
sel_ids = [str(pid) for pid in selected_ids_in_order]
|
|
pos_map: Dict[str, int] = {pid: idx for idx, pid in enumerate(sel_ids)}
|
|
selected_set = set(sel_ids)
|
|
|
|
# Find current minimum number so we can place selected before it
|
|
try:
|
|
min_num = min((it[0] for it in q.queue), default=0)
|
|
except Exception:
|
|
min_num = 0
|
|
new_num_start = min_num - len(sel_ids)
|
|
|
|
new_items: List[Tuple] = []
|
|
for old in q.queue:
|
|
pid = str(old[1])
|
|
prompt = old[2]
|
|
# Clean rename metadata for all items to avoid accidental drift
|
|
if isinstance(prompt, dict) and ('workflow' in prompt or 'name' in prompt):
|
|
prompt = {k: v for k, v in prompt.items() if k not in ['workflow', 'name']}
|
|
if pid in selected_set:
|
|
idx = pos_map.get(pid, 0)
|
|
new_num = new_num_start + idx
|
|
new_items.append((new_num, old[1], prompt, old[3], old[4]))
|
|
else:
|
|
# Keep existing number to preserve original order
|
|
new_items.append((old[0], old[1], prompt, old[3], old[4]))
|
|
|
|
q.queue = new_items
|
|
heapq.heapify(q.queue)
|
|
q.server.queue_updated()
|
|
try:
|
|
# Nudge workers to pick up newly promoted items immediately
|
|
q.not_empty.notify_all()
|
|
except Exception:
|
|
pass
|
|
|
|
def _is_prompt_allowed_while_paused(self, prompt_id: str) -> bool:
|
|
try:
|
|
return str(prompt_id) in self._run_selected_remaining
|
|
except Exception:
|
|
return False
|
|
|
|
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)
|
|
|
|
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.thumbs.embed_webp_metadata(img, save_kwargs, workflow_json)
|
|
elif params['image_format'] == 'png':
|
|
self.thumbs.embed_png_metadata(img, save_kwargs, workflow_json)
|
|
|
|
img.save(cache_path, **save_kwargs)
|
|
return cache_path
|
|
|
|
def _count_samplers_from_prompt(self, prompt: Optional[Any]) -> int:
|
|
"""Heuristic: count nodes that look like samplers in a ComfyUI prompt JSON.
|
|
|
|
Supports both dict-of-nodes and list-in-'nodes' formats.
|
|
"""
|
|
try:
|
|
if prompt is None:
|
|
return 0
|
|
nodes: List[Dict[str, Any]] = []
|
|
if isinstance(prompt, dict) and isinstance(prompt.get('nodes'), list):
|
|
nodes = [n for n in prompt.get('nodes') if isinstance(n, dict)]
|
|
elif isinstance(prompt, dict):
|
|
# dict keyed by node_id -> {class_type, ...}
|
|
for v in prompt.values():
|
|
if isinstance(v, dict) and ('class_type' in v or 'class' in v):
|
|
nodes.append(v)
|
|
count = 0
|
|
for n in nodes:
|
|
try:
|
|
ct = str(n.get('class_type') or n.get('class') or '')
|
|
if not ct:
|
|
continue
|
|
# Broad match: any class type containing 'Sampler'
|
|
if 'sampler' in ct.lower():
|
|
count += 1
|
|
except Exception:
|
|
pass
|
|
return count
|
|
except Exception:
|
|
return 0
|
|
|
|
def _get_total_samplers(self, pid: str, prompt: Optional[Any]) -> int:
|
|
"""Return cached or computed total sampler count for a prompt id.
|
|
|
|
Prefer prompt JSON if available; otherwise fall back to DB stored workflow.
|
|
Cache the result to avoid repeated parsing.
|
|
"""
|
|
total = 0
|
|
calculated_from_prompt = False
|
|
|
|
# First, try to calculate from the provided prompt (if available)
|
|
# This ensures we get the most up-to-date count
|
|
if prompt is not None:
|
|
try:
|
|
total = int(self._count_samplers_from_prompt(prompt) or 0)
|
|
calculated_from_prompt = True
|
|
except Exception:
|
|
total = 0
|
|
calculated_from_prompt = False
|
|
|
|
# If we successfully calculated from prompt, always use that value and update cache
|
|
if calculated_from_prompt and total >= 0:
|
|
try:
|
|
self._samplers_total[pid] = int(total)
|
|
except Exception:
|
|
pass
|
|
return int(total)
|
|
|
|
# Only use cached value if we couldn't calculate from prompt
|
|
if not calculated_from_prompt:
|
|
try:
|
|
if pid in self._samplers_total:
|
|
cached_val = int(self._samplers_total.get(pid) or 0)
|
|
if cached_val >= 0:
|
|
return cached_val
|
|
except Exception:
|
|
pass
|
|
|
|
# Fall back to DB stored workflow
|
|
try:
|
|
row = self.db.get_job(pid)
|
|
if row and row.get('workflow'):
|
|
try:
|
|
wf = row['workflow']
|
|
wf_obj = json.loads(wf) if isinstance(wf, str) else wf
|
|
total = int(self._count_samplers_from_prompt(wf_obj) or 0)
|
|
except Exception:
|
|
total = 0
|
|
except Exception:
|
|
total = 0
|
|
|
|
# Cache the result
|
|
try:
|
|
self._samplers_total[pid] = int(total)
|
|
except Exception:
|
|
pass
|
|
return int(total)
|
|
|
|
|