136 lines
4.4 KiB
Python
136 lines
4.4 KiB
Python
"""
|
|
Async helper utilities for ComfyUI-Distributed.
|
|
"""
|
|
import asyncio
|
|
import threading
|
|
import uuid
|
|
import execution
|
|
import server
|
|
from typing import Optional, Any, Coroutine
|
|
from .network import get_server_loop
|
|
|
|
def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -> Any:
|
|
"""
|
|
Run async coroutine in server's event loop and wait for result.
|
|
|
|
This is useful when you need to run async code from a synchronous context
|
|
but want to use the server's existing event loop instead of creating a new one.
|
|
|
|
Args:
|
|
coro: The coroutine to run
|
|
timeout: Optional timeout in seconds
|
|
|
|
Returns:
|
|
The result of the coroutine
|
|
|
|
Raises:
|
|
TimeoutError: If the operation times out
|
|
Exception: Any exception raised by the coroutine
|
|
"""
|
|
event = threading.Event()
|
|
result = None
|
|
error = None
|
|
|
|
async def wrapper():
|
|
nonlocal result, error
|
|
try:
|
|
result = await coro
|
|
except Exception as e:
|
|
error = e
|
|
finally:
|
|
event.set()
|
|
|
|
# Schedule on server's event loop
|
|
loop = get_server_loop()
|
|
asyncio.run_coroutine_threadsafe(wrapper(), loop)
|
|
|
|
# Wait for completion
|
|
if not event.wait(timeout):
|
|
raise TimeoutError(f"Async operation timed out after {timeout} seconds")
|
|
|
|
if error:
|
|
raise error
|
|
return result
|
|
|
|
|
|
prompt_server = server.PromptServer.instance
|
|
|
|
|
|
def _summarize_node_errors(node_errors: dict) -> str:
|
|
if not isinstance(node_errors, dict) or not node_errors:
|
|
return ""
|
|
|
|
parts = []
|
|
for node_id, entry in node_errors.items():
|
|
if not isinstance(entry, dict):
|
|
continue
|
|
class_type = str(entry.get("class_type") or "UnknownNode")
|
|
for err in entry.get("errors", []):
|
|
if not isinstance(err, dict):
|
|
continue
|
|
message = str(err.get("message") or "validation error")
|
|
details = str(err.get("details") or "").strip()
|
|
parts.append(
|
|
f"{class_type}#{node_id}: {message}{f' ({details})' if details else ''}"
|
|
)
|
|
if len(parts) >= 5:
|
|
return " | ".join(parts)
|
|
return " | ".join(parts)
|
|
|
|
|
|
class PromptValidationError(RuntimeError):
|
|
"""Raised when a prompt fails ComfyUI validation with structured context."""
|
|
|
|
def __init__(self, error_payload, node_errors=None):
|
|
payload = error_payload if isinstance(error_payload, dict) else {
|
|
"type": "prompt_validation_failed",
|
|
"message": str(error_payload),
|
|
"details": "",
|
|
"extra_info": {},
|
|
}
|
|
self.validation_error = dict(payload)
|
|
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
|
|
|
if self.node_errors:
|
|
details = str(self.validation_error.get("details") or "").strip()
|
|
if not details:
|
|
summary = _summarize_node_errors(self.node_errors)
|
|
if summary:
|
|
self.validation_error["details"] = summary
|
|
|
|
merged = dict(self.validation_error)
|
|
if self.node_errors:
|
|
merged["node_errors"] = self.node_errors
|
|
super().__init__(f"Invalid prompt: {merged}")
|
|
|
|
|
|
async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
|
|
"""Validate and queue a prompt via ComfyUI's prompt queue."""
|
|
payload = {"prompt": prompt_obj}
|
|
payload = prompt_server.trigger_on_prompt(payload)
|
|
prompt = payload["prompt"]
|
|
|
|
prompt_id = str(uuid.uuid4())
|
|
valid = await execution.validate_prompt(prompt_id, prompt, None)
|
|
if not valid[0]:
|
|
error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
|
|
node_errors = valid[3] if len(valid) > 3 else {}
|
|
raise PromptValidationError(error_payload, node_errors)
|
|
|
|
extra_data = {}
|
|
if workflow_meta:
|
|
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
|
if client_id:
|
|
extra_data["client_id"] = client_id
|
|
|
|
sensitive = {}
|
|
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
|
|
if key in extra_data:
|
|
sensitive[key] = extra_data.pop(key)
|
|
|
|
number = getattr(prompt_server, "number", 0)
|
|
prompt_server.number = number + 1
|
|
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
|
prompt_server.prompt_queue.put(prompt_queue_item)
|
|
return prompt_id
|