471 lines
18 KiB
Python
471 lines
18 KiB
Python
"""Dynamic Modal proxy nodes for ComfyUI execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import threading
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from dataclasses import dataclass
|
|
from typing import Any, Protocol
|
|
|
|
from comfy_api.latest import _io as io
|
|
|
|
from .serialization import deserialize_node_outputs, serialize_node_inputs
|
|
|
|
logger = logging.getLogger(__name__)
|
|
MODAL_MAP_INPUT_NODE_ID = "ModalMapInput"
|
|
_PROXY_CACHE_CONTEXT_ID_KEY = "__comfy_modal_proxy_cache_context_id__"
|
|
_VOLATILE_PROXY_CACHE_KEYS = frozenset(
|
|
{
|
|
"prompt_id",
|
|
"remote_session",
|
|
"clear_remote_session",
|
|
"extra_data",
|
|
"requires_volume_reload",
|
|
"volume_reload_marker",
|
|
"uploaded_volume_paths",
|
|
}
|
|
)
|
|
|
|
|
|
class RemoteExecutorClient(Protocol):
|
|
"""Execution client interface used by Modal proxy nodes."""
|
|
|
|
def execute_payload(self, payload: Mapping[str, Any], kwargs: Mapping[str, Any]) -> Sequence[Any]:
|
|
"""Execute a serialized Modal payload and return its outputs."""
|
|
|
|
async def execute_payload_async(
|
|
self,
|
|
payload: Mapping[str, Any],
|
|
kwargs: Mapping[str, Any],
|
|
) -> Sequence[Any]:
|
|
"""Execute a serialized Modal payload asynchronously and return its outputs."""
|
|
|
|
|
|
class ModalRemoteExecutorClient:
|
|
"""Default execution client backed by the remote Modal app module."""
|
|
|
|
def execute_payload(self, payload: Mapping[str, Any], kwargs: Mapping[str, Any]) -> Sequence[Any]:
|
|
"""Serialize inputs, invoke the remote engine, and deserialize outputs."""
|
|
from .remote.modal_app import invoke_remote_engine
|
|
|
|
response = invoke_remote_engine(dict(payload), serialize_node_inputs(kwargs))
|
|
return deserialize_node_outputs(response)
|
|
|
|
async def execute_payload_async(
|
|
self,
|
|
payload: Mapping[str, Any],
|
|
kwargs: Mapping[str, Any],
|
|
) -> Sequence[Any]:
|
|
"""Serialize inputs, invoke the remote engine asynchronously, and deserialize outputs."""
|
|
from .remote.modal_app import invoke_remote_engine_async
|
|
|
|
response = await invoke_remote_engine_async(dict(payload), serialize_node_inputs(kwargs))
|
|
return deserialize_node_outputs(response)
|
|
|
|
|
|
_REMOTE_EXECUTOR_CLIENT_FACTORY: Callable[[], RemoteExecutorClient] = ModalRemoteExecutorClient
|
|
_PROXY_NODE_CACHE: dict[str, type[io.ComfyNode]] = {}
|
|
_PROXY_EXECUTION_CONTEXTS_LOCK = threading.Lock()
|
|
_PROXY_EXECUTION_CONTEXTS: dict[str, "_ProxyExecutionContext"] = {}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _ProxyExecutionContext:
|
|
"""Run-scoped execution context used to rehydrate cache-friendly proxy payloads."""
|
|
|
|
execution_payload: dict[str, Any]
|
|
|
|
|
|
def set_remote_executor_client_factory(
|
|
factory: Callable[[], RemoteExecutorClient] | None,
|
|
) -> None:
|
|
"""Install a custom client factory, primarily for tests."""
|
|
global _REMOTE_EXECUTOR_CLIENT_FACTORY
|
|
_REMOTE_EXECUTOR_CLIENT_FACTORY = factory or ModalRemoteExecutorClient
|
|
|
|
|
|
def get_remote_executor_client() -> RemoteExecutorClient:
|
|
"""Instantiate the configured execution client."""
|
|
return _REMOTE_EXECUTOR_CLIENT_FACTORY()
|
|
|
|
|
|
async def _execute_payload_async(
|
|
client: RemoteExecutorClient,
|
|
payload: Mapping[str, Any],
|
|
kwargs: Mapping[str, Any],
|
|
) -> Sequence[Any]:
|
|
"""Execute one Modal payload through the client, adapting sync clients when needed."""
|
|
execute_payload_async = getattr(client, "execute_payload_async", None)
|
|
if callable(execute_payload_async):
|
|
result = execute_payload_async(payload, kwargs)
|
|
if inspect.isawaitable(result):
|
|
return await result
|
|
return result
|
|
|
|
execute_payload = getattr(client, "execute_payload", None)
|
|
if not callable(execute_payload):
|
|
raise TypeError("Remote executor client must define execute_payload or execute_payload_async.")
|
|
return await asyncio.to_thread(execute_payload, payload, kwargs)
|
|
|
|
|
|
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 _normalize_proxy_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
|
|
"""Convert ComfyUI INPUT_IS_LIST proxy kwargs back into ordinary runtime values."""
|
|
normalized_kwargs: dict[str, Any] = {}
|
|
for input_name, input_value in kwargs.items():
|
|
if isinstance(input_value, list) and len(input_value) == 1:
|
|
normalized_kwargs[str(input_name)] = input_value[0]
|
|
continue
|
|
normalized_kwargs[str(input_name)] = input_value
|
|
return normalized_kwargs
|
|
|
|
|
|
def _normalize_proxy_payload(payload: Any) -> Mapping[str, Any]:
|
|
"""Convert ComfyUI INPUT_IS_LIST payload wrappers back into one payload mapping."""
|
|
if isinstance(payload, list) and len(payload) == 1:
|
|
payload = payload[0]
|
|
if isinstance(payload, str):
|
|
payload = json.loads(payload)
|
|
if not isinstance(payload, Mapping):
|
|
raise TypeError("original_node_data must be a mapping or JSON object.")
|
|
return payload
|
|
|
|
|
|
def _normalize_prompt_id(value: Any) -> str | None:
|
|
"""Return one non-empty prompt id string when available."""
|
|
if value is None:
|
|
return None
|
|
prompt_id = str(value).strip()
|
|
return prompt_id or None
|
|
|
|
|
|
def _payload_is_local_cache_safe(payload: Mapping[str, Any]) -> bool:
|
|
"""Return whether one proxy payload can safely reuse local ComfyUI outputs across prompt runs."""
|
|
split_proxy_payloads = payload.get("split_proxy_payloads")
|
|
if isinstance(split_proxy_payloads, Mapping):
|
|
return all(
|
|
isinstance(nested_payload, Mapping) and _payload_is_local_cache_safe(nested_payload)
|
|
for nested_payload in split_proxy_payloads.values()
|
|
)
|
|
if isinstance(split_proxy_payloads, Sequence) and not isinstance(split_proxy_payloads, (str, bytes, bytearray)):
|
|
return all(
|
|
isinstance(nested_payload, Mapping) and _payload_is_local_cache_safe(nested_payload)
|
|
for nested_payload in split_proxy_payloads
|
|
)
|
|
|
|
for phase_name in ("static_phase", "mapped_phase"):
|
|
phase_payload = payload.get(phase_name)
|
|
if isinstance(phase_payload, Mapping) and not _payload_is_local_cache_safe(phase_payload):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _sanitize_cache_surface_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
|
|
"""Strip run-scoped fields from one proxy payload before exposing it to ComfyUI caching."""
|
|
sanitized_payload = dict(payload)
|
|
for field_name in _VOLATILE_PROXY_CACHE_KEYS:
|
|
sanitized_payload.pop(field_name, None)
|
|
|
|
split_proxy_payloads = sanitized_payload.get("split_proxy_payloads")
|
|
if isinstance(split_proxy_payloads, Mapping):
|
|
sanitized_payload["split_proxy_payloads"] = {
|
|
str(phase_name): _sanitize_cache_surface_payload(dict(phase_payload))
|
|
for phase_name, phase_payload in split_proxy_payloads.items()
|
|
if isinstance(phase_payload, Mapping)
|
|
}
|
|
elif isinstance(split_proxy_payloads, Sequence) and not isinstance(
|
|
split_proxy_payloads,
|
|
(str, bytes, bytearray),
|
|
):
|
|
sanitized_payload["split_proxy_payloads"] = [
|
|
_sanitize_cache_surface_payload(dict(phase_payload))
|
|
for phase_payload in split_proxy_payloads
|
|
if isinstance(phase_payload, Mapping)
|
|
]
|
|
|
|
for phase_name in ("static_phase", "mapped_phase"):
|
|
phase_payload = sanitized_payload.get(phase_name)
|
|
if isinstance(phase_payload, Mapping):
|
|
sanitized_payload[phase_name] = _sanitize_cache_surface_payload(dict(phase_payload))
|
|
return sanitized_payload
|
|
|
|
|
|
def register_cache_friendly_proxy_payload(
|
|
node_id: str,
|
|
payload: Mapping[str, Any],
|
|
) -> Mapping[str, Any]:
|
|
"""Return the payload that should be embedded in the proxy node input for local cache reuse."""
|
|
if not _payload_is_local_cache_safe(payload):
|
|
return dict(payload)
|
|
|
|
sanitized_payload = _sanitize_cache_surface_payload(payload)
|
|
sanitized_payload[_PROXY_CACHE_CONTEXT_ID_KEY] = str(node_id)
|
|
with _PROXY_EXECUTION_CONTEXTS_LOCK:
|
|
_PROXY_EXECUTION_CONTEXTS[str(node_id)] = _ProxyExecutionContext(
|
|
execution_payload=dict(payload),
|
|
)
|
|
logger.debug(
|
|
"Registered cache-friendly Modal proxy payload for node_id=%s prompt_id=%s session_backed=%s.",
|
|
node_id,
|
|
_normalize_prompt_id(payload.get("prompt_id")),
|
|
payload.get("remote_session") is not None,
|
|
)
|
|
return sanitized_payload
|
|
|
|
|
|
def _rehydrate_proxy_payload(
|
|
payload: Mapping[str, Any],
|
|
*,
|
|
unique_id: str | None,
|
|
) -> Mapping[str, Any]:
|
|
"""Restore any execution-scoped fields stripped from a cache-friendly proxy payload."""
|
|
context_id = unique_id
|
|
if context_id is None:
|
|
candidate_context_id = payload.get(_PROXY_CACHE_CONTEXT_ID_KEY)
|
|
if candidate_context_id is not None:
|
|
normalized_context_id = str(candidate_context_id).strip()
|
|
context_id = normalized_context_id or None
|
|
if context_id is None:
|
|
return payload
|
|
|
|
with _PROXY_EXECUTION_CONTEXTS_LOCK:
|
|
context = _PROXY_EXECUTION_CONTEXTS.get(str(context_id))
|
|
if context is None:
|
|
return payload
|
|
|
|
return dict(context.execution_payload)
|
|
|
|
|
|
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,
|
|
) -> 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)
|
|
]
|
|
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],
|
|
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(kwargs.pop(io.Hidden.unique_id.name, None))
|
|
payload = _rehydrate_proxy_payload(
|
|
_normalize_proxy_payload(kwargs.pop(payload_input_name, None)),
|
|
unique_id=unique_id,
|
|
)
|
|
|
|
outputs = tuple(
|
|
await _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),
|
|
)
|
|
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:
|
|
nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] = _PROXY_NODE_CACHE[proxy_node_id]
|
|
nodes_module.NODE_DISPLAY_NAME_MAPPINGS[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,
|
|
)
|
|
nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] = proxy_class
|
|
nodes_module.NODE_DISPLAY_NAME_MAPPINGS[proxy_node_id] = "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,
|
|
) -> 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 + (str(is_output_node),),
|
|
)
|
|
|
|
if proxy_node_id in _PROXY_NODE_CACHE:
|
|
nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] = _PROXY_NODE_CACHE[proxy_node_id]
|
|
nodes_module.NODE_DISPLAY_NAME_MAPPINGS[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,
|
|
)
|
|
nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] = proxy_class
|
|
nodes_module.NODE_DISPLAY_NAME_MAPPINGS[proxy_node_id] = "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
|
|
|
|
|
|
class ModalUniversalExecutor(io.ComfyNode):
|
|
"""Base debug node for Modal execution routing."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Define a minimal dev-only schema for the base proxy node."""
|
|
return io.Schema(
|
|
node_id="ModalUniversalExecutor",
|
|
display_name="Modal Universal Executor",
|
|
category="Modal",
|
|
description=(
|
|
"Debug entrypoint for Modal-backed execution. Production rewrites use "
|
|
"signature-preserving dynamic proxy variants."
|
|
),
|
|
inputs=[io.AnyType.Input("original_node_data")],
|
|
outputs=[io.AnyType.Output(display_name="output")],
|
|
accept_all_inputs=True,
|
|
is_dev_only=True,
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, original_node_data: Any, **kwargs: Any) -> io.NodeOutput:
|
|
"""Execute the base debug proxy node through the configured client."""
|
|
if isinstance(original_node_data, str):
|
|
original_node_data = json.loads(original_node_data)
|
|
outputs = tuple(get_remote_executor_client().execute_node(original_node_data, kwargs))
|
|
return io.NodeOutput(*outputs)
|
|
|
|
|
|
class ModalMapInput(io.ComfyNode):
|
|
"""Queue-time marker node that turns one remote boundary input into mapped parallel work."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""Expose a simple any-to-any adapter node for mapped remote execution."""
|
|
return io.Schema(
|
|
node_id=MODAL_MAP_INPUT_NODE_ID,
|
|
display_name="Modal Map Input",
|
|
category="Modal",
|
|
description=(
|
|
"Pass-through marker for data-parallel Modal execution. "
|
|
"When used inside a remote-marked component, list inputs and batched tensors "
|
|
"can fan out across multiple Modal executions and reassemble automatically."
|
|
),
|
|
inputs=[io.AnyType.Input("value")],
|
|
outputs=[io.AnyType.Output(display_name="value")],
|
|
is_experimental=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, value: Any) -> io.NodeOutput:
|
|
"""Pass the input value through unchanged at runtime."""
|
|
return io.NodeOutput(value)
|