335 lines
12 KiB
Python
335 lines
12 KiB
Python
"""Dynamic ComfyUI proxy-node construction and mapping registration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from comfy_api.latest import _io as io
|
|
|
|
if __package__:
|
|
from .proxy_payloads import (
|
|
_normalize_prompt_id,
|
|
_normalize_proxy_kwargs,
|
|
_normalize_proxy_payload,
|
|
_normalize_scheduler_list_outputs,
|
|
_pop_proxy_hidden_value,
|
|
_prompt_id_from_extra_pnginfo,
|
|
_rehydrate_proxy_payload,
|
|
_unwrap_proxy_singleton,
|
|
)
|
|
from .remote_executor_router import get_remote_executor_client
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from proxy_payloads import (
|
|
_normalize_prompt_id,
|
|
_normalize_proxy_kwargs,
|
|
_normalize_proxy_payload,
|
|
_normalize_scheduler_list_outputs,
|
|
_pop_proxy_hidden_value,
|
|
_prompt_id_from_extra_pnginfo,
|
|
_rehydrate_proxy_payload,
|
|
_unwrap_proxy_singleton,
|
|
)
|
|
from remote_executor_router import get_remote_executor_client
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MODAL_ARTIFACT_FINALIZER_NODE_ID = "ModalArtifactFinalizer"
|
|
MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID = "ModalParallelLocalPassthrough"
|
|
MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID = "ModalLocalBridgeMaterializer"
|
|
MODAL_COMPONENT_COMPLETION_OUTPUT_NAME = "modal_component_complete"
|
|
_PROXY_NODE_CACHE: dict[str, type[io.ComfyNode]] = {}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProxyNodeFactoryHooks:
|
|
"""Entrypoint-owned execution and static-node dependencies."""
|
|
|
|
modal_workflow_execution_slot: Callable[..., Any]
|
|
signal_parallel_local_dispatch: Callable[[Mapping[str, Any]], bool]
|
|
execute_payload_async: Callable[..., Any]
|
|
modal_universal_executor: type[io.ComfyNode]
|
|
modal_artifact_finalizer: type[io.ComfyNode]
|
|
modal_parallel_local_passthrough: type[io.ComfyNode]
|
|
modal_local_bridge_materializer: type[io.ComfyNode]
|
|
|
|
|
|
_PROXY_NODE_FACTORY_HOOKS: ProxyNodeFactoryHooks | None = None
|
|
|
|
|
|
def configure_proxy_node_factory_hooks(hooks: ProxyNodeFactoryHooks) -> None:
|
|
"""Install residual executor and static-node dependencies."""
|
|
global _PROXY_NODE_FACTORY_HOOKS
|
|
_PROXY_NODE_FACTORY_HOOKS = hooks
|
|
|
|
|
|
def _factory_hooks() -> ProxyNodeFactoryHooks:
|
|
"""Return configured factory hooks or fail on invalid import order."""
|
|
if _PROXY_NODE_FACTORY_HOOKS is None:
|
|
raise RuntimeError("Proxy node factory hooks have not been configured.")
|
|
return _PROXY_NODE_FACTORY_HOOKS
|
|
|
|
def _output_spec(io_type: str, name: str, is_list: bool) -> io.Output:
|
|
"""Create a v3 output specification from a legacy ComfyUI return type."""
|
|
comfy_type = io.AnyType if io_type == "*" else io.Custom(io_type)
|
|
return comfy_type.Output(display_name=name, is_output_list=is_list)
|
|
|
|
|
|
|
|
def _normalized_output_metadata(
|
|
original_class: type[Any],
|
|
) -> tuple[tuple[str, ...], tuple[str, ...], tuple[bool, ...]]:
|
|
"""Normalize output metadata from a source node class."""
|
|
if hasattr(original_class, "GET_SCHEMA"):
|
|
original_class.GET_SCHEMA()
|
|
|
|
output_types = tuple(getattr(original_class, "RETURN_TYPES", ("*",))) or ("*",)
|
|
default_names = tuple(f"output_{index}" for index, _ in enumerate(output_types))
|
|
output_names = tuple(getattr(original_class, "RETURN_NAMES", default_names))
|
|
output_is_list = tuple(getattr(original_class, "OUTPUT_IS_LIST", (False,) * len(output_types)))
|
|
|
|
if len(output_names) < len(output_types):
|
|
output_names = output_names + default_names[len(output_names) :]
|
|
if len(output_is_list) < len(output_types):
|
|
output_is_list = output_is_list + (False,) * (len(output_types) - len(output_is_list))
|
|
|
|
return output_types, output_names[: len(output_types)], output_is_list[: len(output_types)]
|
|
|
|
|
|
def _proxy_node_id(original_class_type: str, output_types: Sequence[str]) -> str:
|
|
"""Build a stable proxy node identifier for an original node signature."""
|
|
digest = hashlib.sha256(
|
|
json.dumps({"class_type": original_class_type, "outputs": list(output_types)}).encode(
|
|
"utf-8"
|
|
)
|
|
).hexdigest()[:12]
|
|
return f"ModalUniversalExecutor_{digest}"
|
|
|
|
|
|
def _build_proxy_node_class(
|
|
node_id: str,
|
|
proxy_display_name: str,
|
|
payload_input_name: str,
|
|
output_types: tuple[str, ...],
|
|
output_names: tuple[str, ...],
|
|
output_is_list: tuple[bool, ...],
|
|
*,
|
|
is_output_node: bool,
|
|
include_completion_output: bool = False,
|
|
) -> type[io.ComfyNode]:
|
|
"""Create a v3 proxy node that mirrors an original node output signature."""
|
|
|
|
class _DynamicModalExecutor(io.ComfyNode):
|
|
"""Internal proxy node that forwards execution to Modal."""
|
|
|
|
OUTPUT_NODE = is_output_node
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Return a schema that accepts any original node inputs."""
|
|
outputs = [
|
|
_output_spec(io_type, name, is_list)
|
|
for io_type, name, is_list in zip(output_types, output_names, output_is_list, strict=False)
|
|
]
|
|
if include_completion_output:
|
|
outputs.append(
|
|
io.Boolean.Output(
|
|
display_name=MODAL_COMPONENT_COMPLETION_OUTPUT_NAME,
|
|
)
|
|
)
|
|
return io.Schema(
|
|
node_id=node_id,
|
|
display_name=proxy_display_name,
|
|
category="Modal",
|
|
description=(
|
|
"Internal proxy node that forwards a rewritten Modal execution "
|
|
"payload to a Modal-backed runtime."
|
|
),
|
|
inputs=[
|
|
io.AnyType.Input(payload_input_name),
|
|
],
|
|
outputs=outputs,
|
|
is_input_list=True,
|
|
accept_all_inputs=True,
|
|
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo],
|
|
is_dev_only=True,
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(cls, **kwargs: Any) -> io.NodeOutput:
|
|
"""Forward the execution payload to the configured remote executor."""
|
|
unique_id = _normalize_prompt_id(
|
|
_unwrap_proxy_singleton(
|
|
_pop_proxy_hidden_value(cls, kwargs, io.Hidden.unique_id)
|
|
)
|
|
)
|
|
prompt_id = _prompt_id_from_extra_pnginfo(
|
|
_pop_proxy_hidden_value(cls, kwargs, io.Hidden.extra_pnginfo)
|
|
)
|
|
payload = _rehydrate_proxy_payload(
|
|
_normalize_proxy_payload(kwargs.pop(payload_input_name, None)),
|
|
unique_id=unique_id,
|
|
prompt_id=prompt_id,
|
|
)
|
|
|
|
async with _factory_hooks().modal_workflow_execution_slot(payload):
|
|
_factory_hooks().signal_parallel_local_dispatch(payload)
|
|
outputs = _normalize_scheduler_list_outputs(
|
|
payload,
|
|
await _factory_hooks().execute_payload_async(
|
|
get_remote_executor_client(),
|
|
payload,
|
|
_normalize_proxy_kwargs(kwargs),
|
|
),
|
|
)
|
|
logger.debug(
|
|
"Remote execution completed for payload kind=%s with %d outputs.",
|
|
payload.get("payload_kind"),
|
|
len(outputs),
|
|
)
|
|
if include_completion_output:
|
|
return io.NodeOutput(*outputs, True)
|
|
return io.NodeOutput(*outputs)
|
|
|
|
_DynamicModalExecutor.__name__ = f"DynamicModalExecutor_{node_id}"
|
|
return _DynamicModalExecutor
|
|
|
|
|
|
def ensure_modal_proxy_node_registered(
|
|
original_class_type: str,
|
|
original_class: type[Any],
|
|
nodes_module: Any,
|
|
) -> str:
|
|
"""Register and return a proxy node id for the supplied original node class."""
|
|
output_types, output_names, output_is_list = _normalized_output_metadata(original_class)
|
|
proxy_node_id = _proxy_node_id(original_class_type, output_types)
|
|
|
|
if proxy_node_id in _PROXY_NODE_CACHE:
|
|
_register_modal_node(
|
|
nodes_module,
|
|
proxy_node_id,
|
|
_PROXY_NODE_CACHE[proxy_node_id],
|
|
"Modal Universal Executor",
|
|
)
|
|
return proxy_node_id
|
|
|
|
proxy_class = _build_proxy_node_class(
|
|
node_id=proxy_node_id,
|
|
proxy_display_name="Modal Universal Executor",
|
|
payload_input_name="original_node_data",
|
|
output_types=output_types,
|
|
output_names=output_names,
|
|
output_is_list=output_is_list,
|
|
is_output_node=False,
|
|
)
|
|
_register_modal_node(
|
|
nodes_module,
|
|
proxy_node_id,
|
|
proxy_class,
|
|
"Modal Universal Executor",
|
|
)
|
|
_PROXY_NODE_CACHE[proxy_node_id] = proxy_class
|
|
logger.info("Registered Modal proxy node %s for %s", proxy_node_id, original_class_type)
|
|
return proxy_node_id
|
|
|
|
|
|
def ensure_modal_component_proxy_node_registered(
|
|
output_types: Sequence[str],
|
|
output_names: Sequence[str],
|
|
output_is_list: Sequence[bool],
|
|
nodes_module: Any,
|
|
*,
|
|
is_output_node: bool,
|
|
include_completion_output: bool = False,
|
|
) -> str:
|
|
"""Register and return a proxy node id for a remote component signature."""
|
|
normalized_output_types = tuple(output_types)
|
|
normalized_output_names = tuple(output_names)
|
|
normalized_output_is_list = tuple(output_is_list)
|
|
proxy_node_id = _proxy_node_id(
|
|
"ModalRemoteComponent",
|
|
normalized_output_types
|
|
+ tuple(f"name:{name}" for name in normalized_output_names)
|
|
+ tuple(f"list:{is_list}" for is_list in normalized_output_is_list)
|
|
+ (str(is_output_node), str(include_completion_output)),
|
|
)
|
|
|
|
if proxy_node_id in _PROXY_NODE_CACHE:
|
|
_register_modal_node(
|
|
nodes_module,
|
|
proxy_node_id,
|
|
_PROXY_NODE_CACHE[proxy_node_id],
|
|
"Modal Remote Component",
|
|
)
|
|
return proxy_node_id
|
|
|
|
proxy_class = _build_proxy_node_class(
|
|
node_id=proxy_node_id,
|
|
proxy_display_name="Modal Remote Component",
|
|
payload_input_name="original_node_data",
|
|
output_types=normalized_output_types,
|
|
output_names=normalized_output_names,
|
|
output_is_list=normalized_output_is_list,
|
|
is_output_node=is_output_node,
|
|
include_completion_output=include_completion_output,
|
|
)
|
|
_register_modal_node(
|
|
nodes_module,
|
|
proxy_node_id,
|
|
proxy_class,
|
|
"Modal Remote Component",
|
|
)
|
|
_PROXY_NODE_CACHE[proxy_node_id] = proxy_class
|
|
logger.info("Registered Modal component proxy node %s", proxy_node_id)
|
|
return proxy_node_id
|
|
|
|
|
|
def _register_modal_node(
|
|
nodes_module: Any,
|
|
node_id: str,
|
|
node_class: type[io.ComfyNode],
|
|
display_name: str,
|
|
) -> None:
|
|
"""Register a runtime Modal node with the plugin's ComfyUI module identity."""
|
|
node_class.RELATIVE_PYTHON_MODULE = _factory_hooks().modal_universal_executor.RELATIVE_PYTHON_MODULE
|
|
nodes_module.NODE_CLASS_MAPPINGS[node_id] = node_class
|
|
nodes_module.NODE_DISPLAY_NAME_MAPPINGS[node_id] = display_name
|
|
|
|
|
|
def ensure_modal_artifact_finalizer_registered(nodes_module: Any) -> None:
|
|
"""Register the internal artifact-finalization sink in a ComfyUI node mapping."""
|
|
_register_modal_node(
|
|
nodes_module,
|
|
MODAL_ARTIFACT_FINALIZER_NODE_ID,
|
|
_factory_hooks().modal_artifact_finalizer,
|
|
"Modal Artifact Finalizer",
|
|
)
|
|
|
|
|
|
def ensure_modal_parallel_local_passthrough_registered(nodes_module: Any) -> None:
|
|
"""Register the internal parallel local-branch dispatch gate."""
|
|
_register_modal_node(
|
|
nodes_module,
|
|
MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID,
|
|
_factory_hooks().modal_parallel_local_passthrough,
|
|
"Modal Parallel Local Passthrough",
|
|
)
|
|
|
|
|
|
def ensure_modal_local_bridge_materializer_registered(nodes_module: Any) -> None:
|
|
"""Register the internal durable-bridge local materializer."""
|
|
_register_modal_node(
|
|
nodes_module,
|
|
MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID,
|
|
_factory_hooks().modal_local_bridge_materializer,
|
|
"Modal Local Bridge Materializer",
|
|
)
|
|
|
|
|
|
|