189 lines
6.6 KiB
Python
189 lines
6.6 KiB
Python
"""
|
|
Async helper utilities for ComfyUI-Distributed.
|
|
"""
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
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}")
|
|
|
|
|
|
def validate_execution_targets(prompt, targets):
|
|
"""Reject malformed explicit roots instead of falling back to all outputs."""
|
|
if (not isinstance(targets, list) or not targets
|
|
or any(not isinstance(node_id, str) or not node_id or node_id not in prompt
|
|
for node_id in targets)
|
|
or len(set(targets)) != len(targets)):
|
|
raise PromptValidationError({
|
|
"type": "invalid_partial_execution_targets",
|
|
"message": "partial_execution_targets must be a non-empty list of unique existing node IDs",
|
|
"details": "", "extra_info": {},
|
|
})
|
|
return list(targets)
|
|
|
|
|
|
async def queue_prompt_payload(
|
|
prompt_obj,
|
|
workflow_meta=None,
|
|
client_id=None,
|
|
include_queue_metadata=False,
|
|
*,
|
|
partial_execution_targets=None,
|
|
):
|
|
"""Validate and queue a prompt, preserving any explicitly requested roots."""
|
|
targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
|
|
if partial_execution_targets is not None else None)
|
|
payload = {"prompt": prompt_obj}
|
|
if targets is not None:
|
|
payload["partial_execution_targets"] = list(targets)
|
|
payload = prompt_server.trigger_on_prompt(payload)
|
|
prompt = payload["prompt"]
|
|
if targets is not None:
|
|
# Hooks may replace the envelope or mutate its target list. The caller's
|
|
# requested roots remain authoritative, but must still exist afterwards.
|
|
validate_execution_targets(prompt, targets)
|
|
|
|
prompt_id = str(uuid.uuid4())
|
|
valid = await execution.validate_prompt(prompt_id, prompt, targets)
|
|
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)
|
|
|
|
if targets is not None:
|
|
node_errors = valid[3] if len(valid) > 3 else {}
|
|
target_errors = {
|
|
node_id: entry for node_id, entry in node_errors.items()
|
|
if node_id in targets or (isinstance(entry, dict)
|
|
and set(entry.get("dependent_outputs", [])) & set(targets))
|
|
}
|
|
if set(valid[2]) != set(targets) or target_errors:
|
|
raise PromptValidationError({
|
|
"type": "partial_execution_targets_failed_validation",
|
|
"message": "Not all requested execution targets passed validation",
|
|
"details": "Refusing to queue a subset or additional output roots",
|
|
"extra_info": {"requested_targets": targets},
|
|
}, node_errors)
|
|
|
|
extra_data = {"create_time": int(time.time() * 1000)}
|
|
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)
|
|
|
|
if include_queue_metadata:
|
|
return {
|
|
"prompt_id": prompt_id,
|
|
"number": number,
|
|
"node_errors": {},
|
|
}
|
|
|
|
return prompt_id
|