mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-07-27 16:26:39 +00:00
46c8311d14
This commit adds support for node execution functions defined as async. When a node's execution function is defined as async, we can continue executing other nodes while it is processing. Standard uses of `await` should "just work", but people will still have to be careful if they spawn actual threads. Because torch doesn't really have async/await versions of functions, this won't particularly help with most locally-executing nodes, but it does work for e.g. web requests to other machines. In addition to the execute function, the `VALIDATE_INPUTS` and `check_lazy_status` functions can also be defined as async, though we'll only resolve one node at a time right now for those.
289 lines
10 KiB
Python
289 lines
10 KiB
Python
from typing import TypedDict, Dict, Optional
|
|
from typing_extensions import override
|
|
from PIL import Image
|
|
from enum import Enum
|
|
from abc import ABC
|
|
from tqdm import tqdm
|
|
from comfy_execution.graph import DynamicPrompt
|
|
from protocol import BinaryEventTypes
|
|
|
|
class NodeState(Enum):
|
|
Pending = "pending"
|
|
Running = "running"
|
|
Finished = "finished"
|
|
Error = "error"
|
|
|
|
class NodeProgressState(TypedDict):
|
|
"""
|
|
A class to represent the state of a node's progress.
|
|
"""
|
|
state: NodeState
|
|
value: float
|
|
max: float
|
|
|
|
class ProgressHandler(ABC):
|
|
"""
|
|
Abstract base class for progress handlers.
|
|
Progress handlers receive progress updates and display them in various ways.
|
|
"""
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
self.enabled = True
|
|
|
|
def set_registry(self, registry: "ProgressRegistry"):
|
|
pass
|
|
|
|
def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
|
|
"""Called when a node starts processing"""
|
|
pass
|
|
|
|
def update_handler(self, node_id: str, value: float, max_value: float,
|
|
state: NodeProgressState, prompt_id: str, image: Optional[Image.Image] = None):
|
|
"""Called when a node's progress is updated"""
|
|
pass
|
|
|
|
def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
|
|
"""Called when a node finishes processing"""
|
|
pass
|
|
|
|
def reset(self):
|
|
"""Called when the progress registry is reset"""
|
|
pass
|
|
|
|
def enable(self):
|
|
"""Enable this handler"""
|
|
self.enabled = True
|
|
|
|
def disable(self):
|
|
"""Disable this handler"""
|
|
self.enabled = False
|
|
|
|
class CLIProgressHandler(ProgressHandler):
|
|
"""
|
|
Handler that displays progress using tqdm progress bars in the CLI.
|
|
"""
|
|
def __init__(self):
|
|
super().__init__("cli")
|
|
self.progress_bars: Dict[str, tqdm] = {}
|
|
|
|
@override
|
|
def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
|
|
# Create a new tqdm progress bar
|
|
if node_id not in self.progress_bars:
|
|
self.progress_bars[node_id] = tqdm(
|
|
total=state["max"],
|
|
desc=f"Node {node_id}",
|
|
unit="steps",
|
|
leave=True,
|
|
position=len(self.progress_bars)
|
|
)
|
|
|
|
@override
|
|
def update_handler(self, node_id: str, value: float, max_value: float,
|
|
state: NodeProgressState, prompt_id: str, image: Optional[Image.Image] = None):
|
|
# Handle case where start_handler wasn't called
|
|
if node_id not in self.progress_bars:
|
|
self.progress_bars[node_id] = tqdm(
|
|
total=max_value,
|
|
desc=f"Node {node_id}",
|
|
unit="steps",
|
|
leave=True,
|
|
position=len(self.progress_bars)
|
|
)
|
|
self.progress_bars[node_id].update(value)
|
|
else:
|
|
# Update existing progress bar
|
|
if max_value != self.progress_bars[node_id].total:
|
|
self.progress_bars[node_id].total = max_value
|
|
# Calculate the update amount (difference from current position)
|
|
current_position = self.progress_bars[node_id].n
|
|
update_amount = value - current_position
|
|
if update_amount > 0:
|
|
self.progress_bars[node_id].update(update_amount)
|
|
|
|
@override
|
|
def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
|
|
# Complete and close the progress bar if it exists
|
|
if node_id in self.progress_bars:
|
|
# Ensure the bar shows 100% completion
|
|
remaining = state["max"] - self.progress_bars[node_id].n
|
|
if remaining > 0:
|
|
self.progress_bars[node_id].update(remaining)
|
|
self.progress_bars[node_id].close()
|
|
del self.progress_bars[node_id]
|
|
|
|
@override
|
|
def reset(self):
|
|
# Close all progress bars
|
|
for bar in self.progress_bars.values():
|
|
bar.close()
|
|
self.progress_bars.clear()
|
|
|
|
class WebUIProgressHandler(ProgressHandler):
|
|
"""
|
|
Handler that sends progress updates to the WebUI via WebSockets.
|
|
"""
|
|
def __init__(self, server_instance):
|
|
super().__init__("webui")
|
|
self.server_instance = server_instance
|
|
|
|
def set_registry(self, registry: "ProgressRegistry"):
|
|
self.registry = registry
|
|
|
|
def _send_progress_state(self, prompt_id: str, nodes: Dict[str, NodeProgressState]):
|
|
"""Send the current progress state to the client"""
|
|
if self.server_instance is None:
|
|
return
|
|
|
|
# Only send info for non-pending nodes
|
|
active_nodes = {
|
|
node_id: {
|
|
"value": state["value"],
|
|
"max": state["max"],
|
|
"state": state["state"].value,
|
|
"node_id": node_id,
|
|
"prompt_id": prompt_id,
|
|
"display_node_id": self.registry.dynprompt.get_display_node_id(node_id),
|
|
"parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
|
|
"real_node_id": self.registry.dynprompt.get_real_node_id(node_id)
|
|
}
|
|
for node_id, state in nodes.items()
|
|
if state["state"] != NodeState.Pending
|
|
}
|
|
|
|
# Send a combined progress_state message with all node states
|
|
self.server_instance.send_sync("progress_state", {
|
|
"prompt_id": prompt_id,
|
|
"nodes": active_nodes
|
|
})
|
|
|
|
@override
|
|
def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
|
|
# Send progress state of all nodes
|
|
if self.registry:
|
|
self._send_progress_state(prompt_id, self.registry.nodes)
|
|
|
|
@override
|
|
def update_handler(self, node_id: str, value: float, max_value: float,
|
|
state: NodeProgressState, prompt_id: str, image: Optional[Image.Image] = None):
|
|
# Send progress state of all nodes
|
|
if self.registry:
|
|
self._send_progress_state(prompt_id, self.registry.nodes)
|
|
if image:
|
|
metadata = {
|
|
"node_id": node_id,
|
|
"prompt_id": prompt_id,
|
|
"display_node_id": self.registry.dynprompt.get_display_node_id(node_id),
|
|
"parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
|
|
"real_node_id": self.registry.dynprompt.get_real_node_id(node_id)
|
|
}
|
|
self.server_instance.send_sync(BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA, (image, metadata), self.server_instance.client_id)
|
|
|
|
|
|
@override
|
|
def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
|
|
# Send progress state of all nodes
|
|
if self.registry:
|
|
self._send_progress_state(prompt_id, self.registry.nodes)
|
|
|
|
class ProgressRegistry:
|
|
"""
|
|
Registry that maintains node progress state and notifies registered handlers.
|
|
"""
|
|
def __init__(self, prompt_id: str, dynprompt: DynamicPrompt):
|
|
self.prompt_id = prompt_id
|
|
self.dynprompt = dynprompt
|
|
self.nodes: Dict[str, NodeProgressState] = {}
|
|
self.handlers: Dict[str, ProgressHandler] = {}
|
|
|
|
def register_handler(self, handler: ProgressHandler) -> None:
|
|
"""Register a progress handler"""
|
|
self.handlers[handler.name] = handler
|
|
|
|
def unregister_handler(self, handler_name: str) -> None:
|
|
"""Unregister a progress handler"""
|
|
if handler_name in self.handlers:
|
|
# Allow handler to clean up resources
|
|
self.handlers[handler_name].reset()
|
|
del self.handlers[handler_name]
|
|
|
|
def enable_handler(self, handler_name: str) -> None:
|
|
"""Enable a progress handler"""
|
|
if handler_name in self.handlers:
|
|
self.handlers[handler_name].enable()
|
|
|
|
def disable_handler(self, handler_name: str) -> None:
|
|
"""Disable a progress handler"""
|
|
if handler_name in self.handlers:
|
|
self.handlers[handler_name].disable()
|
|
|
|
def ensure_entry(self, node_id: str) -> NodeProgressState:
|
|
"""Ensure a node entry exists"""
|
|
if node_id not in self.nodes:
|
|
self.nodes[node_id] = NodeProgressState(
|
|
state = NodeState.Pending,
|
|
value = 0,
|
|
max = 1
|
|
)
|
|
return self.nodes[node_id]
|
|
|
|
def start_progress(self, node_id: str) -> None:
|
|
"""Start progress tracking for a node"""
|
|
entry = self.ensure_entry(node_id)
|
|
entry["state"] = NodeState.Running
|
|
entry["value"] = 0.0
|
|
entry["max"] = 1.0
|
|
|
|
# Notify all enabled handlers
|
|
for handler in self.handlers.values():
|
|
if handler.enabled:
|
|
handler.start_handler(node_id, entry, self.prompt_id)
|
|
|
|
def update_progress(self, node_id: str, value: float, max_value: float, image: Optional[Image.Image]) -> None:
|
|
"""Update progress for a node"""
|
|
entry = self.ensure_entry(node_id)
|
|
entry["state"] = NodeState.Running
|
|
entry["value"] = value
|
|
entry["max"] = max_value
|
|
|
|
# Notify all enabled handlers
|
|
for handler in self.handlers.values():
|
|
if handler.enabled:
|
|
handler.update_handler(node_id, value, max_value, entry, self.prompt_id, image)
|
|
|
|
def finish_progress(self, node_id: str) -> None:
|
|
"""Finish progress tracking for a node"""
|
|
entry = self.ensure_entry(node_id)
|
|
entry["state"] = NodeState.Finished
|
|
entry["value"] = entry["max"]
|
|
|
|
# Notify all enabled handlers
|
|
for handler in self.handlers.values():
|
|
if handler.enabled:
|
|
handler.finish_handler(node_id, entry, self.prompt_id)
|
|
|
|
def reset_handlers(self) -> None:
|
|
"""Reset all handlers"""
|
|
for handler in self.handlers.values():
|
|
handler.reset()
|
|
|
|
# Global registry instance
|
|
global_progress_registry: ProgressRegistry = ProgressRegistry(prompt_id="", dynprompt=DynamicPrompt({}))
|
|
|
|
def reset_progress_state(prompt_id: str, dynprompt: DynamicPrompt) -> None:
|
|
global global_progress_registry
|
|
|
|
# Reset existing handlers if registry exists
|
|
if global_progress_registry is not None:
|
|
global_progress_registry.reset_handlers()
|
|
|
|
# Create new registry
|
|
global_progress_registry = ProgressRegistry(prompt_id, dynprompt)
|
|
|
|
def add_progress_handler(handler: ProgressHandler) -> None:
|
|
handler.set_registry(global_progress_registry)
|
|
global_progress_registry.register_handler(handler)
|
|
|
|
def get_progress_state() -> ProgressRegistry:
|
|
return global_progress_registry
|