521 lines
20 KiB
Python
521 lines
20 KiB
Python
"""Headless ComfyUI prompt-server shims and execution event tracing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from io import BytesIO
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Any, Callable
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CloudPromptServerHooks:
|
|
"""Callbacks supplied by the cloud orchestration and execution owners."""
|
|
|
|
collapse_cache_slot: Callable[[Any, bool], Any]
|
|
emit_cloud_info: Callable[..., None]
|
|
meaningful_progress_values: Callable[[dict[str, Any]], tuple[float, float] | None]
|
|
|
|
|
|
_PROMPT_SERVER_HOOKS: CloudPromptServerHooks | None = None
|
|
|
|
|
|
def configure_cloud_prompt_server_hooks(hooks: CloudPromptServerHooks) -> None:
|
|
"""Install prompt-server callbacks without importing upward into the entrypoint."""
|
|
global _PROMPT_SERVER_HOOKS
|
|
_PROMPT_SERVER_HOOKS = hooks
|
|
|
|
|
|
def _prompt_server_hooks() -> CloudPromptServerHooks:
|
|
"""Return configured callbacks or fail with a clear import-order error."""
|
|
if _PROMPT_SERVER_HOOKS is None:
|
|
raise RuntimeError("Cloud prompt-server hooks have not been configured.")
|
|
return _PROMPT_SERVER_HOOKS
|
|
|
|
|
|
def _collapse_cache_slot(slot_values: Any, is_list: bool) -> Any:
|
|
"""Delegate cache-slot normalization to prompt execution."""
|
|
return _prompt_server_hooks().collapse_cache_slot(slot_values, is_list)
|
|
|
|
|
|
def _emit_cloud_info(message: str, *args: Any) -> None:
|
|
"""Delegate timestamped cloud logging to the stable entrypoint."""
|
|
_prompt_server_hooks().emit_cloud_info(message, *args)
|
|
|
|
|
|
def _meaningful_progress_values(
|
|
node_state: dict[str, Any],
|
|
) -> tuple[float, float] | None:
|
|
"""Delegate progress-state filtering to the stable entrypoint."""
|
|
return _prompt_server_hooks().meaningful_progress_values(node_state)
|
|
|
|
|
|
class _NullPromptServer:
|
|
"""Minimal PromptExecutor server stub for headless subgraph execution."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the no-op prompt server state."""
|
|
self.client_id: str | None = None
|
|
self.last_node_id: str | None = None
|
|
self.last_prompt_id: str | None = None
|
|
|
|
def send_sync(
|
|
self, event: str, data: dict[str, Any], client_id: str | None
|
|
) -> None:
|
|
"""Discard PromptExecutor progress and status events."""
|
|
logger.debug(
|
|
"Suppressed remote prompt event %s for client %s.", event, client_id
|
|
)
|
|
|
|
|
|
class _HeadlessPromptQueue:
|
|
"""In-memory PromptQueue compatibility surface for remote custom-node hooks."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the queue collections exposed by ComfyUI's PromptQueue."""
|
|
self._lock = threading.RLock()
|
|
self.queue: list[Any] = []
|
|
self.currently_running: dict[int, Any] = {}
|
|
self.history: dict[str, Any] = {}
|
|
self.flags: dict[str, Any] = {}
|
|
|
|
def put(self, item: Any) -> None:
|
|
"""Retain a custom-node requeue request for compatibility diagnostics."""
|
|
with self._lock:
|
|
self.queue.append(item)
|
|
logger.warning(
|
|
"A remote custom node queued headless follow-up work; the request is retained "
|
|
"for compatibility but has no background queue consumer."
|
|
)
|
|
|
|
def get_current_queue(self) -> tuple[list[Any], list[Any]]:
|
|
"""Return snapshots of running and pending compatibility queue items."""
|
|
with self._lock:
|
|
return list(self.currently_running.values()), list(self.queue)
|
|
|
|
def get_current_queue_volatile(self) -> tuple[list[Any], list[Any]]:
|
|
"""Return shallow running and pending queue snapshots."""
|
|
return self.get_current_queue()
|
|
|
|
def get_tasks_remaining(self) -> int:
|
|
"""Return the number of running and pending compatibility tasks."""
|
|
with self._lock:
|
|
return len(self.currently_running) + len(self.queue)
|
|
|
|
def set_flag(self, name: str, data: Any) -> None:
|
|
"""Store a queue flag for custom nodes that inspect PromptQueue state."""
|
|
with self._lock:
|
|
self.flags[name] = data
|
|
|
|
def get_flags(self, reset: bool = True) -> dict[str, Any]:
|
|
"""Return queue flags and optionally clear them."""
|
|
with self._lock:
|
|
flags = dict(self.flags)
|
|
if reset:
|
|
self.flags.clear()
|
|
return flags
|
|
|
|
|
|
class _HeadlessPromptServerInstance:
|
|
"""Minimal PromptServer.instance replacement for custom-node import side effects."""
|
|
|
|
def __init__(self, node_replace_manager: Any | None) -> None:
|
|
"""Initialize route registration and no-op websocket state."""
|
|
from aiohttp import web
|
|
|
|
self.routes = web.RouteTableDef()
|
|
self.app = web.Application()
|
|
self.node_replace_manager = node_replace_manager
|
|
self.supports = ["custom_nodes_from_web"]
|
|
self.client_id: str | None = None
|
|
self.last_node_id: str | None = None
|
|
self.number = 0
|
|
self.prompt_queue = _HeadlessPromptQueue()
|
|
self.on_prompt_handlers: list[Any] = []
|
|
|
|
def queue_updated(self) -> None:
|
|
"""Accept PromptQueue update notifications without a websocket client."""
|
|
logger.debug("Suppressed headless remote prompt queue update.")
|
|
|
|
async def send(
|
|
self, event: str, data: dict[str, Any], sid: str | None = None
|
|
) -> None:
|
|
"""Discard async websocket sends from import-time custom-node helpers."""
|
|
logger.debug(
|
|
"Suppressed headless remote prompt event %s for client %s.", event, sid
|
|
)
|
|
|
|
def send_sync(
|
|
self, event: str, data: dict[str, Any], sid: str | None = None
|
|
) -> None:
|
|
"""Discard sync websocket sends from import-time custom-node helpers."""
|
|
logger.debug(
|
|
"Suppressed headless remote prompt event %s for client %s.", event, sid
|
|
)
|
|
|
|
def send_progress_text(
|
|
self,
|
|
text: bytes | bytearray | str,
|
|
node_id: str,
|
|
sid: str | None = None,
|
|
) -> None:
|
|
"""Discard node progress text emitted by headless remote node execution."""
|
|
del text
|
|
logger.debug(
|
|
"Suppressed headless remote prompt text for node %s client %s.",
|
|
node_id,
|
|
sid,
|
|
)
|
|
|
|
def add_on_prompt_handler(self, handler: Any) -> None:
|
|
"""Record prompt handlers registered by custom nodes during import."""
|
|
self.on_prompt_handlers.append(handler)
|
|
|
|
|
|
class _TracingPromptServer(_NullPromptServer):
|
|
"""PromptExecutor server stub that records coarse per-node execution timings."""
|
|
|
|
def __init__(
|
|
self,
|
|
prompt_id: str,
|
|
prompt: dict[str, Any],
|
|
status_callback: Callable[[dict[str, Any]], None] | None = None,
|
|
) -> None:
|
|
"""Initialize timing state for a specific prompt execution."""
|
|
super().__init__()
|
|
self.prompt_id = prompt_id
|
|
self.prompt = prompt
|
|
self._status_callback = status_callback
|
|
self._active_node_id: str | None = None
|
|
self._active_node_started_at: float | None = None
|
|
self.last_prompt_id = prompt_id
|
|
self._boundary_outputs_by_node_id: dict[str, list[dict[str, Any]]] = {}
|
|
self._lookup_cache_entry: Callable[[str], Any | None] | None = None
|
|
self._published_boundary_outputs: set[tuple[str, int]] = set()
|
|
|
|
def _classify_node_role(self, class_type: str) -> str:
|
|
"""Return a coarse role name for a node class."""
|
|
normalized = class_type.lower()
|
|
if "loader" in normalized or normalized in {"clipvisionencode"}:
|
|
return "model_load"
|
|
if "ksampler" in normalized or "sampler" in normalized:
|
|
return "sampling"
|
|
if "encode" in normalized:
|
|
return "conditioning"
|
|
return "node"
|
|
|
|
def _log_node_finish(self, reason: str) -> None:
|
|
"""Emit a timing line for the currently active node when one is running."""
|
|
if self._active_node_id is None or self._active_node_started_at is None:
|
|
return
|
|
|
|
node_id = self._active_node_id
|
|
node_info = self.prompt.get(node_id, {})
|
|
class_type = str(node_info.get("class_type", "<unknown>"))
|
|
role = self._classify_node_role(class_type)
|
|
elapsed_seconds = time.perf_counter() - self._active_node_started_at
|
|
_emit_cloud_info(
|
|
"Remote node %s class_type=%s role=%s finished in %.3fs reason=%s",
|
|
node_id,
|
|
class_type,
|
|
role,
|
|
elapsed_seconds,
|
|
reason,
|
|
)
|
|
self._active_node_id = None
|
|
self._active_node_started_at = None
|
|
|
|
def emit_preview_update(
|
|
self,
|
|
*,
|
|
node_id: str,
|
|
preview_image: Any,
|
|
) -> None:
|
|
"""Publish one preview image update through the status callback."""
|
|
if self._status_callback is None:
|
|
return
|
|
|
|
try:
|
|
image_type, image, max_size = preview_image
|
|
image_buffer = BytesIO()
|
|
save_kwargs: dict[str, Any] = {"format": image_type}
|
|
if image_type == "JPEG":
|
|
save_kwargs["quality"] = 95
|
|
elif image_type == "PNG":
|
|
save_kwargs["compress_level"] = 1
|
|
image.save(image_buffer, **save_kwargs)
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to serialize remote preview image for node %s.", node_id
|
|
)
|
|
return
|
|
|
|
try:
|
|
from comfy_execution.progress import get_progress_state
|
|
|
|
registry = get_progress_state()
|
|
display_node_id = registry.dynprompt.get_display_node_id(node_id)
|
|
parent_node_id = registry.dynprompt.get_parent_node_id(node_id)
|
|
real_node_id = registry.dynprompt.get_real_node_id(node_id)
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to resolve preview metadata for remote node %s.", node_id
|
|
)
|
|
display_node_id = node_id
|
|
parent_node_id = None
|
|
real_node_id = node_id
|
|
|
|
self._status_callback(
|
|
{
|
|
"event_type": "preview",
|
|
"node_id": str(node_id),
|
|
"display_node_id": (
|
|
str(display_node_id) if display_node_id is not None else None
|
|
),
|
|
"parent_node_id": (
|
|
str(parent_node_id) if parent_node_id is not None else None
|
|
),
|
|
"real_node_id": str(real_node_id) if real_node_id is not None else None,
|
|
"image_type": str(image_type),
|
|
"image_bytes": image_buffer.getvalue(),
|
|
"max_size": int(max_size) if max_size is not None else None,
|
|
}
|
|
)
|
|
|
|
def configure_boundary_output_stream(
|
|
self,
|
|
*,
|
|
boundary_outputs: list[dict[str, Any]],
|
|
lookup_cache_entry: Callable[[str], Any | None],
|
|
) -> None:
|
|
"""Configure streamed remote boundary-output publication for this execution."""
|
|
outputs_by_node_id: dict[str, list[dict[str, Any]]] = {}
|
|
for boundary_output in boundary_outputs:
|
|
preview_target_node_ids = [
|
|
str(node_id)
|
|
for node_id in boundary_output.get("preview_target_node_ids", [])
|
|
if str(node_id)
|
|
]
|
|
if not preview_target_node_ids:
|
|
continue
|
|
if str(boundary_output.get("io_type")) != "IMAGE":
|
|
continue
|
|
node_id = str(boundary_output["node_id"])
|
|
outputs_by_node_id.setdefault(node_id, []).append(boundary_output)
|
|
|
|
self._boundary_outputs_by_node_id = outputs_by_node_id
|
|
self._lookup_cache_entry = lookup_cache_entry
|
|
self._published_boundary_outputs.clear()
|
|
|
|
def _emit_boundary_outputs_for_node(self, node_id: str | None) -> None:
|
|
"""Publish configured boundary image outputs for one completed node once."""
|
|
if (
|
|
node_id is None
|
|
or self._status_callback is None
|
|
or self._lookup_cache_entry is None
|
|
):
|
|
return
|
|
|
|
boundary_outputs = self._boundary_outputs_by_node_id.get(str(node_id), [])
|
|
if not boundary_outputs:
|
|
return
|
|
|
|
cache_entry = self._lookup_cache_entry(str(node_id))
|
|
if cache_entry is None:
|
|
return
|
|
|
|
cache_outputs = getattr(cache_entry, "outputs", None)
|
|
if not isinstance(cache_outputs, (list, tuple)):
|
|
return
|
|
|
|
for boundary_output in boundary_outputs:
|
|
output_index = int(boundary_output["output_index"])
|
|
publication_key = (str(node_id), output_index)
|
|
if publication_key in self._published_boundary_outputs:
|
|
continue
|
|
if output_index >= len(cache_outputs):
|
|
continue
|
|
|
|
preview_target_node_ids = [
|
|
str(target_node_id)
|
|
for target_node_id in boundary_output.get("preview_target_node_ids", [])
|
|
if str(target_node_id)
|
|
]
|
|
if not preview_target_node_ids:
|
|
continue
|
|
|
|
self._status_callback(
|
|
{
|
|
"event_type": "boundary_output",
|
|
"node_id": str(node_id),
|
|
"output_index": output_index,
|
|
"io_type": str(boundary_output.get("io_type", "")),
|
|
"is_list": bool(boundary_output.get("is_list", False)),
|
|
"preview_target_node_ids": preview_target_node_ids,
|
|
"value": _collapse_cache_slot(
|
|
slot_values=cache_outputs[output_index],
|
|
is_list=bool(boundary_output.get("is_list", False)),
|
|
),
|
|
}
|
|
)
|
|
self._published_boundary_outputs.add(publication_key)
|
|
|
|
def send_sync(
|
|
self, event: str, data: dict[str, Any], client_id: str | None
|
|
) -> None:
|
|
"""Track per-node timing transitions from PromptExecutor progress events."""
|
|
if event == "modal_llm_progress":
|
|
if self._status_callback is None:
|
|
return
|
|
node_id = str(data.get("node_id") or self._active_node_id or "")
|
|
if not node_id:
|
|
return
|
|
progress_event = {
|
|
"event_type": "node_progress",
|
|
"node_id": node_id,
|
|
"display_node_id": node_id,
|
|
"real_node_id": node_id,
|
|
"value": float(data.get("value", 0.0)),
|
|
"max": float(data.get("max", 1.0)),
|
|
"stage": str(data.get("stage") or ""),
|
|
"message": str(data.get("message") or ""),
|
|
"indeterminate": bool(data.get("indeterminate", False)),
|
|
}
|
|
for field_name in (
|
|
"unit",
|
|
"elapsed_seconds",
|
|
"time_to_first_token_seconds",
|
|
"tokens_per_second",
|
|
):
|
|
if data.get(field_name) is not None:
|
|
progress_event[field_name] = data[field_name]
|
|
self._status_callback(progress_event)
|
|
return
|
|
|
|
if event == "executing":
|
|
next_node_id = data.get("node")
|
|
if next_node_id != self._active_node_id:
|
|
self._emit_boundary_outputs_for_node(self._active_node_id)
|
|
self._log_node_finish(reason="next_node")
|
|
if next_node_id is not None and next_node_id != self._active_node_id:
|
|
node_info = self.prompt.get(str(next_node_id), {})
|
|
class_type = str(node_info.get("class_type", "<unknown>"))
|
|
role = self._classify_node_role(class_type)
|
|
if self._status_callback is not None:
|
|
self._status_callback(
|
|
{
|
|
"phase": "executing",
|
|
"active_node_id": str(next_node_id),
|
|
"active_node_class_type": class_type,
|
|
"active_node_role": role,
|
|
}
|
|
)
|
|
self._active_node_id = str(next_node_id)
|
|
self._active_node_started_at = time.perf_counter()
|
|
self.last_node_id = self._active_node_id
|
|
_emit_cloud_info(
|
|
"Remote node %s class_type=%s role=%s started",
|
|
self._active_node_id,
|
|
class_type,
|
|
role,
|
|
)
|
|
return
|
|
|
|
if event == "progress_state":
|
|
if self._status_callback is None:
|
|
return
|
|
|
|
nodes_payload = data.get("nodes")
|
|
if not isinstance(nodes_payload, dict):
|
|
return
|
|
|
|
tracked_node_id = self._active_node_id
|
|
tracked_node_state: dict[str, Any] | None = None
|
|
if tracked_node_id is not None:
|
|
candidate_state = nodes_payload.get(tracked_node_id)
|
|
if isinstance(candidate_state, dict):
|
|
tracked_node_state = candidate_state
|
|
|
|
if tracked_node_state is None:
|
|
for node_state in nodes_payload.values():
|
|
if (
|
|
isinstance(node_state, dict)
|
|
and node_state.get("state") == "running"
|
|
):
|
|
tracked_node_state = node_state
|
|
break
|
|
|
|
if tracked_node_state is None:
|
|
return
|
|
|
|
display_node_id = tracked_node_state.get("display_node_id")
|
|
real_node_id = tracked_node_state.get("real_node_id")
|
|
reported_node_id = (
|
|
display_node_id or real_node_id or tracked_node_state.get("node_id")
|
|
)
|
|
if reported_node_id is None:
|
|
return
|
|
progress_values = _meaningful_progress_values(tracked_node_state)
|
|
if progress_values is None:
|
|
logger.debug(
|
|
"Ignoring non-meaningful remote progress_state for node_id=%s state=%s value=%r max=%r.",
|
|
reported_node_id,
|
|
tracked_node_state.get("state"),
|
|
tracked_node_state.get("value"),
|
|
tracked_node_state.get("max"),
|
|
)
|
|
return
|
|
progress_value, max_value = progress_values
|
|
|
|
self._status_callback(
|
|
{
|
|
"event_type": "node_progress",
|
|
"node_id": str(reported_node_id),
|
|
"display_node_id": (
|
|
str(display_node_id) if display_node_id is not None else None
|
|
),
|
|
"real_node_id": str(real_node_id)
|
|
if real_node_id is not None
|
|
else None,
|
|
"value": progress_value,
|
|
"max": max_value,
|
|
}
|
|
)
|
|
return
|
|
|
|
if event == "executed":
|
|
executed_node_id = data.get("node")
|
|
if (
|
|
executed_node_id is not None
|
|
and str(executed_node_id) == self._active_node_id
|
|
):
|
|
self._log_node_finish(reason="executed")
|
|
if self._status_callback is not None and data.get("output") is not None:
|
|
self._status_callback(
|
|
{
|
|
"event_type": "executed",
|
|
"node_id": str(data.get("node")),
|
|
"display_node_id": (
|
|
str(data["display_node"])
|
|
if data.get("display_node") is not None
|
|
else None
|
|
),
|
|
"output": data.get("output"),
|
|
}
|
|
)
|
|
return
|
|
|
|
if event in {"execution_error", "execution_interrupted", "execution_success"}:
|
|
self._emit_boundary_outputs_for_node(self._active_node_id)
|
|
self._log_node_finish(reason=event)
|
|
if self._status_callback is not None:
|
|
self._status_callback({"phase": event})
|
|
return
|
|
|
|
super().send_sync(event, data, client_id)
|
|
|