From cf55bd58ae1cf33bf84c362af4e6c74c9a394f4e Mon Sep 17 00:00:00 2001 From: Kosinkadink Date: Mon, 11 May 2026 12:31:05 -0700 Subject: [PATCH] Initial RNP/1 client implementation with Tier 3 features MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ComfyUI custom-node package that fetches node descriptors from an RNP/1 (Remote Node Protocol v1) server and exposes them as native V3 nodes inside ComfyUI. Includes: - Manifest + capability handshake, ETag-conditional object_info fetch - Sync execute_remote and async _run_async_polling flows via comfy_api_nodes.util.client.poll_op_raw - Image / mask / audio / video envelope encode + decode - Presigned upload helper for >8 MiB inputs - Tier 3: * TASK_LOST detection + idempotent resume — survives an RNP server restart mid-task * SERVER_BUSY / MAINTENANCE 503 envelopes treated as rate limits with Retry-After honoring and friendly per-second progress text * Out-of-band url_fetch for URI-only output envelopes per the descriptor's remote.url_fetch policy * Per-poll auth refresh so long-running tasks pick up rotated tokens Pairs with Comfy-Org/comfy-rnp-server (private) as the reference server implementation. --- .gitignore | 12 + README.md | 62 +++ __init__.py | 15 + client.py | 431 +++++++++++++++++ protocol.py | 172 +++++++ proxy_node.py | 1179 ++++++++++++++++++++++++++++++++++++++++++++++ registry.py | 177 +++++++ serialization.py | 318 +++++++++++++ 8 files changed, 2366 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 client.py create mode 100644 protocol.py create mode 100644 proxy_node.py create mode 100644 registry.py create mode 100644 serialization.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..55ade59 --- /dev/null +++ b/.gitignore @@ -0,0 +1,12 @@ +__pycache__/ +*.pyc +*.pyo +*.pyd +.venv/ +.env +*.egg-info/ +build/ +dist/ +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..89b5bf1 --- /dev/null +++ b/README.md @@ -0,0 +1,62 @@ +# comfy-remote-nodes + +ComfyUI custom-node package that fetches node descriptors from an +**RNP/1** (Remote Node Protocol v1) server and exposes them as native +V3 nodes inside ComfyUI. + +The reference server lives at +[`Comfy-Org/comfy-rnp-server`](https://github.com/Comfy-Org/comfy-rnp-server) +(private). This client is intentionally provider-agnostic: any server +that speaks RNP/1 can publish nodes through it. + +## Install + +Clone into your ComfyUI `custom_nodes/` directory: + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/Comfy-Org/comfy-remote-nodes.git +``` + +(The directory name on disk doesn't matter — ComfyUI loads the package +via `spec_from_file_location`, so `comfy-remote-nodes/` is fine.) + +## Configure + +Point the extension at one RNP server URL via env var before starting +ComfyUI: + +```bash +export RNP_SERVER_URL=http://127.0.0.1:9190 +python main.py +``` + +If `RNP_SERVER_URL` is unset, the extension registers zero nodes and +logs a single line — ComfyUI starts normally with no remote nodes. + +## What works today + +* Manifest + capability handshake +* `object_info` fetch and dynamic V3 class generation (with ETag + conditional GETs) +* Sync `execute` and async `execute_async` + poll/cancel +* Image / mask / audio / video envelope encode + decode +* Presigned upload helper for >8 MiB inputs +* **Tier 3** (recent): + * `TASK_LOST` detection with idempotent resubmission — survives an + RNP server restart mid-task without losing the workflow + * `SERVER_BUSY` / `MAINTENANCE` 503 envelopes treated as rate + limits with `Retry-After` honoring and friendly per-second + progress text ("Server busy, retrying in 30s...") + * Out-of-band `url_fetch` for URI-only output envelopes, driven by + the descriptor's `remote.url_fetch` policy + * Per-poll auth refresh so long-running tasks can rotate tokens + +## What does not (yet) + +* `validate_inputs` / `fingerprint_inputs` / `check_lazy_status` — + the server returns `501 NOT_IMPLEMENTED`; the client falls through + to ComfyUI's defaults +* Multi-server config file (`comfy_remote_nodes.yaml` or similar) — + today's setup is one server per ComfyUI process via env var +* Streaming execute (server-sent events for incremental previews) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ed15ef7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,15 @@ +"""comfy_remote_nodes — ComfyUI client for the Remote Node Protocol (RNP/1). + +* On boot, fetches ``GET /rnp/v1/manifest`` then + ``GET /rnp/v1/object_info`` from the server URL set via + ``RNP_SERVER_URL``. +* Builds dynamic V3 ``IO.ComfyNode`` subclasses for each descriptor + with a ``remote.endpoints.execute`` entry and registers them with + ComfyUI through the standard ``ComfyExtension`` interface. +* Only ``execute`` is wired today; ``validate_inputs``, + ``fingerprint_inputs`` and ``check_lazy_status`` fall through to + ComfyUI's defaults. +""" +from __future__ import annotations + +from .registry import comfy_entrypoint # noqa: F401 — exposed to ComfyUI diff --git a/client.py b/client.py new file mode 100644 index 0000000..a7895e9 --- /dev/null +++ b/client.py @@ -0,0 +1,431 @@ +"""HTTP client for the RNP/1 server. + +Builds on ``comfy_api_nodes.util.client.sync_op_raw`` so RNP requests +get the same retry / cancellation / friendly-error treatment as local +partner-node calls. Adds the RNP-specific bits ``sync_op_raw`` doesn't +know about: protocol/capability headers, ``If-None-Match`` 304 handling +on ``/object_info``, and the two-phase upload helper (the GCS PUT runs +on plain aiohttp because presigned URLs bypass api.comfy.org). +""" +from __future__ import annotations + +import json +import logging +from typing import Any +from urllib.parse import urljoin + +import aiohttp + +from comfy_api_nodes.util.client import ApiEndpoint, sync_op_raw + +from .protocol import ( + Capability, + Header, + PROTOCOL_VERSION, + RnpProtocolError, +) + +log = logging.getLogger("comfy_remote_nodes.client") + + +# Capability tokens advertised on every outbound request. Currently +# covers the input/output encodings the registered providers (LTXV, +# Gemini) need, plus the ``schema:v3`` marker so descriptors come back +# in the V3 ``get_v1_info()`` shape. +CLIENT_CAPABILITIES = [ + Capability.SCHEMA_V3, + Capability.VIDEO_MP4_INLINE, + Capability.IMAGE_PNG_BASE64, + Capability.MASK_PNG_BASE64, +] + +# Reported in the X-RNP-Client-Version header. +CLIENT_VERSION = "0.2.0" + + +def _default_headers(extra: dict[str, str] | None = None) -> dict[str, str]: + headers = { + Header.PROTOCOL_VERSION: PROTOCOL_VERSION, + Header.CLIENT_VERSION: CLIENT_VERSION, + Header.CLIENT_CAPABILITIES: json.dumps(CLIENT_CAPABILITIES), + } + if extra: + headers.update(extra) + return headers + + +async def fetch_manifest(server_url: str, *, timeout_s: float = 30.0) -> dict[str, Any]: + """``GET /rnp/v1/manifest`` — protocol + capability handshake.""" + body = await sync_op_raw( + None, + ApiEndpoint("rnp/v1/manifest", "GET", headers=_default_headers()), + timeout=timeout_s, + as_binary=False, + max_retries=2, + wait_label="Fetching manifest", + final_label_on_success=None, + monitor_progress=False, + base_url=server_url, + auth_headers={}, + error_parser=_parse_rnp_error, + is_rate_limited=_is_rnp_backpressure, + ) + if not isinstance(body, dict): + raise RnpProtocolError( + f"Expected JSON manifest, got {type(body).__name__}", + ) + return body + + +async def fetch_object_info( + server_url: str, + *, + timeout_s: float = 30.0, + etag: str | None = None, +) -> tuple[dict[str, Any] | None, str | None, int | None]: + """``GET /rnp/v1/object_info`` — node descriptor map. + + Returns ``(body, etag, max_age_s)``. When the server replies ``304 + Not Modified`` (``If-None-Match`` matched), ``body`` is ``None`` so + the caller can keep its cached descriptor map. ``max_age_s`` is the + Cache-Control hint (or None if absent). + """ + headers = _default_headers() + if etag: + headers["If-None-Match"] = etag + + captured: dict[str, Any] = {"etag": None, "max_age": None} + + def _capture(resp_headers: dict[str, str]) -> None: + captured["etag"] = resp_headers.get("etag") + captured["max_age"] = _parse_max_age(resp_headers.get("cache-control") or "") + + body = await sync_op_raw( + None, + ApiEndpoint("rnp/v1/object_info", "GET", headers=headers), + timeout=timeout_s, + as_binary=False, + max_retries=2, + wait_label="Fetching object_info", + final_label_on_success=None, + monitor_progress=False, + base_url=server_url, + auth_headers={}, + allow_304=True, + response_header_validator=_capture, + error_parser=_parse_rnp_error, + is_rate_limited=_is_rnp_backpressure, + ) + if body is None: # 304 Not Modified + return None, captured["etag"], captured["max_age"] + if not isinstance(body, dict): + raise RnpProtocolError( + f"Expected JSON object_info, got {type(body).__name__}", + ) + return body, captured["etag"], captured["max_age"] + + +def _parse_max_age(cache_control: str) -> int | None: + """Extract ``max-age=`` from a Cache-Control header. + + Returns ``None`` when the directive is missing or malformed; the + caller treats that as 'no caching hint'. + """ + for part in cache_control.split(","): + part = part.strip().lower() + if part.startswith("max-age="): + try: + return int(part.split("=", 1)[1]) + except (IndexError, ValueError): + return None + return None + + +async def upload_asset( + server_url: str, + data: bytes, + *, + content_type: str | None = None, + timeout_s: float = 120.0, + auth_headers: dict[str, str] | None = None, +) -> str: + """Stage ``data`` on the RNP server and return its download URL. + + Two-phase: ``POST /rnp/v1/uploads`` to allocate the slot and get + back ``{upload_url, download_url}``, then ``PUT upload_url`` with + the raw bytes. ``auth_headers`` are forwarded on the allocate call + so the server can delegate to api.comfy.org/customers/storage with + the user's credentials. The PUT runs on plain aiohttp because the + presigned URL bypasses api.comfy.org. + """ + meta_body = await sync_op_raw( + None, + ApiEndpoint( + "rnp/v1/uploads", "POST", + headers=_default_headers(), + ), + data={"content_type": content_type or "application/octet-stream"}, + timeout=timeout_s, + as_binary=False, + max_retries=2, + wait_label="Allocating upload slot", + final_label_on_success=None, + monitor_progress=False, + base_url=server_url, + auth_headers=auth_headers or {}, + error_parser=_parse_rnp_error, + is_rate_limited=_is_rnp_backpressure, + ) + if not isinstance(meta_body, dict): + raise RnpProtocolError( + f"Expected JSON upload response, got {type(meta_body).__name__}", + ) + upload_url = meta_body.get("upload_url") + download_url = meta_body.get("download_url") + if not isinstance(upload_url, str) or not isinstance(download_url, str): + raise RnpProtocolError( + f"upload response missing upload_url/download_url: {meta_body!r}", + ) + + put_headers = {"Content-Type": content_type or "application/octet-stream"} + timeout = aiohttp.ClientTimeout(total=timeout_s) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.put(upload_url, data=data, headers=put_headers) as put_resp: + if put_resp.status >= 400: + raise RnpProtocolError( + f"PUT to upload_url failed: HTTP {put_resp.status}", + status=put_resp.status, + ) + return download_url + + +async def execute_remote( + server_url: str, + schema_hash: str | None, + endpoint_path: str, + inputs: dict[str, Any], + context: dict[str, Any], + *, + timeout_s: float = 600.0, + extra_headers: dict[str, str] | None = None, + node_cls: type | None = None, + rate_limit_label: Any = None, +) -> dict[str, Any]: + """``POST /rnp/v1/nodes//execute`` — run the remote node. + + Returns the parsed JSON body (typically ``{"outputs": [...]}``). + Raises ``RnpProtocolError`` on non-2xx responses with the + structured-error fields populated. ``node_cls`` is forwarded to + ``sync_op_raw`` so progress/interrupt checks key off the right + node id; ``None`` is fine for callers without a workflow node. + """ + headers = _default_headers() + if schema_hash: + headers[Header.SCHEMA_HASH] = schema_hash + if extra_headers: + headers.update(extra_headers) + body = await sync_op_raw( + node_cls, + ApiEndpoint(endpoint_path.lstrip("/"), "POST", headers=headers), + data={"inputs": inputs, "context": context}, + timeout=timeout_s, + as_binary=False, + max_retries=2, + wait_label="Processing", + final_label_on_success=None, + monitor_progress=False, + base_url=server_url, + auth_headers=extra_headers or {}, + error_parser=_parse_rnp_error, + is_rate_limited=_is_rnp_backpressure, + rate_limit_label=rate_limit_label or _rnp_rate_limit_label, + ) + if not isinstance(body, dict): + raise RnpProtocolError( + f"Expected JSON execute response, got {type(body).__name__}", + ) + return body + + +async def execute_remote_async( + server_url: str, + schema_hash: str | None, + node_id: str, + inputs: dict[str, Any], + context: dict[str, Any], + *, + timeout_s: float = 60.0, + extra_headers: dict[str, str] | None = None, + idempotency_key: str | None = None, + node_cls: type | None = None, + rate_limit_label: Any = None, +) -> dict[str, Any]: + """``POST /rnp/v1/nodes//execute_async`` — kick off the + background runner. + + Returns ``{"task_id": "...", "poll_interval": }``. The caller + drives ``poll_op`` until a terminal status. ``idempotency_key`` is + forwarded as ``X-RNP-Idempotency-Key`` so a network-retry of this + POST returns the same task_id rather than spawning a duplicate run. + """ + headers = _default_headers() + if schema_hash: + headers[Header.SCHEMA_HASH] = schema_hash + if idempotency_key: + headers[Header.IDEMPOTENCY_KEY] = idempotency_key + if extra_headers: + headers.update(extra_headers) + body = await sync_op_raw( + node_cls, + ApiEndpoint( + f"rnp/v1/nodes/{node_id}/execute_async", "POST", headers=headers, + ), + data={"inputs": inputs, "context": context}, + timeout=timeout_s, + as_binary=False, + max_retries=2, + wait_label="Submitting", + final_label_on_success=None, + monitor_progress=False, + base_url=server_url, + auth_headers=extra_headers or {}, + error_parser=_parse_rnp_error, + is_rate_limited=_is_rnp_backpressure, + ) + if not isinstance(body, dict) or "task_id" not in body: + raise RnpProtocolError( + f"execute_async response missing task_id: {body!r}", + ) + return body + + +async def poll_op( + server_url: str, + task_id: str, + *, + timeout_s: float = 30.0, + node_cls: type | None = None, +) -> dict[str, Any]: + """``GET /rnp/v1/tasks/`` — single poll-response read. + + Returns the canonical poll-response dict (see + ``comfy_rnp_protocol.tasks``). Caller schedules the loop and the + sleep between polls. + """ + body = await sync_op_raw( + node_cls, + ApiEndpoint( + f"rnp/v1/tasks/{task_id}", "GET", headers=_default_headers(), + ), + timeout=timeout_s, + as_binary=False, + max_retries=2, + wait_label=None, + final_label_on_success=None, + monitor_progress=False, + base_url=server_url, + auth_headers={}, + error_parser=_parse_rnp_error, + is_rate_limited=_is_rnp_backpressure, + ) + if not isinstance(body, dict) or "status" not in body: + raise RnpProtocolError( + f"poll response missing status: {body!r}", + ) + return body + + +async def cancel_op( + server_url: str, + task_id: str, + *, + timeout_s: float = 10.0, +) -> None: + """``POST /rnp/v1/tasks//cancel`` — best-effort cancel. + + Failures are swallowed: if the task already finished or the server + is unreachable, the user has already pressed cancel and we don't + want to surface a second exception over their original interrupt. + """ + try: + await sync_op_raw( + None, + ApiEndpoint( + f"rnp/v1/tasks/{task_id}/cancel", "POST", + headers=_default_headers(), + ), + timeout=timeout_s, + as_binary=False, + max_retries=0, + wait_label=None, + final_label_on_success=None, + monitor_progress=False, + base_url=server_url, + auth_headers={}, + error_parser=_parse_rnp_error, + ) + except Exception as e: # noqa: BLE001 + log.warning("RNP cancel %s failed (best-effort): %s", task_id, e) + + +def _parse_rnp_error(status: int, body: Any) -> Exception | None: + """``sync_op_raw`` error_parser hook. + + Surfaces structured RNP error envelopes (``{"error": {...}}``) as + ``RnpProtocolError`` so the proxy_node can map ``code`` / + ``details.upstream_status`` to the literal frontend dialog strings. + Returns ``None`` for unstructured 4xx/5xx so the default + ``_friendly_http_message`` path keeps firing — that already + produces the magic 401/402 strings for missing-auth / no-credits + failures the server reflects with the real upstream status. + """ + if isinstance(body, dict) and isinstance(body.get("error"), dict): + return RnpProtocolError.from_response_body(body, status) + return None + + +# Wire codes the server emits when it wants the client to wait and +# retry (operator drain / capacity ceiling). Treated identically to a +# 429 by ``sync_op_raw`` so they consume the rate-limit retry counter +# and honor ``Retry-After`` instead of falling through the generic +# 5xx exponential-backoff path. +_BACKPRESSURE_CODES = frozenset({"SERVER_BUSY", "MAINTENANCE"}) + + +def _is_rnp_backpressure(status: int, body: Any) -> bool: + """``sync_op_raw`` is_rate_limited hook for SERVER_BUSY / MAINTENANCE. + + Both arrive as 503 with a structured RNP error envelope; promoting + them to "rate limited" makes the client honor ``Retry-After`` and + use ``max_retries_on_rate_limit`` instead of the generic 5xx + counter, matching the spec's "the server is asking you to wait" + semantics. + """ + if status != 503 or not isinstance(body, dict): + return False + err = body.get("error") + if not isinstance(err, dict): + return False + return err.get("code") in _BACKPRESSURE_CODES + + +def _rnp_rate_limit_label(status: int, body: Any, retry_after_s: float) -> str | None: + """``sync_op_raw`` rate_limit_label hook for SERVER_BUSY / MAINTENANCE. + + Returns the per-second progress label shown while waiting out a + backpressure 503 — distinct copy for capacity vs. operator drain so + users can tell why the workflow is stalled. ``None`` lets the + default ``cfg.wait_label`` ("Waiting for server") stand for plain + 429s and any 503 the RNP server didn't tag with a known code. + """ + secs = max(1, int(round(retry_after_s))) + if isinstance(body, dict): + err = body.get("error") if isinstance(body.get("error"), dict) else None + code = err.get("code") if err else None + if code == "SERVER_BUSY": + return f"Server busy, retrying in {secs}s..." + if code == "MAINTENANCE": + return f"Server in maintenance, back in {secs}s..." + if status == 429: + return f"Rate limited, retrying in {secs}s..." + return None diff --git a/protocol.py b/protocol.py new file mode 100644 index 0000000..73971e0 --- /dev/null +++ b/protocol.py @@ -0,0 +1,172 @@ +"""Vendored copy of the parts of ``comfy_rnp_protocol`` the client needs. + +The full protocol package lives in ``comfy-rnp-server/comfy_rnp_protocol`` +and is the source of truth for these constants and helpers — this file +is a thin mirror so the ComfyUI extension can ship without depending on +the server package being on ``PYTHONPATH``. Keep it byte-compatible with +the upstream module. +""" +from __future__ import annotations + +import base64 +from typing import Any + +PROTOCOL_NAME = "comfy-rnp" +PROTOCOL_VERSION = "1.0" +SCHEMA_HASH_ALGORITHM = "sha256" +DEFAULT_MAX_INLINE_PAYLOAD_BYTES = 8 * 1024 * 1024 +DEFAULT_POLL_INTERVAL_S = 1.5 + +# Defaults for descriptor.remote.execution fields the client must honour. +# Values mirror comfy_api_nodes.util.client.poll_op_raw so RNP-routed +# nodes behave identically to a direct upstream call when the descriptor +# leaves a field unset. +DEFAULT_MAX_POLL_ATTEMPTS = 160 +DEFAULT_TIMEOUT_PER_POLL_S = 120.0 +DEFAULT_MAX_RETRIES_PER_POLL = 10 +DEFAULT_RETRY_DELAY_PER_POLL_S = 1.0 +DEFAULT_RETRY_BACKOFF_PER_POLL = 1.4 +DEFAULT_CANCEL_TIMEOUT_S = 10.0 +DEFAULT_SUBMIT_MAX_RETRIES = 2 +DEFAULT_SUBMIT_RETRY_BACKOFF = 2.0 +DEFAULT_SUBMIT_RETRY_DELAY_S = 1.0 + + +class IdempotencyMode: + CLIENT_KEY = "client_key" + NONE = "none" + + +class ExecutionMode: + REQUEST_RESPONSE = "request_response" + ASYNC_POLLING = "async_polling" + + +class TaskStatus: + PENDING = "pending" + RUNNING = "running" + DONE = "done" + ERROR = "error" + CANCELLED = "cancelled" + + +TERMINAL_TASK_STATES = frozenset({ + TaskStatus.DONE, + TaskStatus.ERROR, + TaskStatus.CANCELLED, +}) + + +class ErrorCode: + INPUT_INVALID = "INPUT_INVALID" + AUTH_FAILED = "AUTH_FAILED" + INSUFFICIENT_CREDITS = "INSUFFICIENT_CREDITS" + RATE_LIMITED = "RATE_LIMITED" + PROVIDER_UNAVAILABLE = "PROVIDER_UNAVAILABLE" + PAYLOAD_TOO_LARGE = "PAYLOAD_TOO_LARGE" + VERSION_INCOMPATIBLE = "VERSION_INCOMPATIBLE" + SCHEMA_HASH_MISMATCH = "SCHEMA_HASH_MISMATCH" + TIMEOUT = "TIMEOUT" + INTERNAL = "INTERNAL" + NOT_FOUND = "NOT_FOUND" + NOT_IMPLEMENTED = "NOT_IMPLEMENTED" + # 404 from a poll/cancel that carried X-RNP-Idempotency-Key — the + # server lost the task (typically a restart). The client replays + # the original submit with the same key to spawn a fresh task. + TASK_LOST = "TASK_LOST" + # Server-side back-pressure / operator drain — both ride 503 with + # Retry-After. SERVER_BUSY = transient saturation, MAINTENANCE = + # operator action. + SERVER_BUSY = "SERVER_BUSY" + MAINTENANCE = "MAINTENANCE" + + +class Header: + PROTOCOL_VERSION = "X-RNP-Protocol-Version" + CLIENT_VERSION = "X-RNP-Client-Version" + CLIENT_CAPABILITIES = "X-RNP-Client-Capabilities" + SCHEMA_HASH = "X-RNP-Schema-Hash" + IDEMPOTENCY_KEY = "X-RNP-Idempotency-Key" + + +class Capability: + IMAGE_PNG_BASE64 = "image:png_base64" + IMAGE_JPEG_BASE64 = "image:jpeg_base64" + IMAGE_RAW_CHW_UINT8 = "image:raw_b64_chw_uint8" + VIDEO_MP4_INLINE = "video:mp4_inline" + AUDIO_MP3_INLINE = "audio:mp3_inline" + MASK_PNG_BASE64 = "mask:png_base64" + SCHEMA_V3 = "schema:v3" + ASYNC_EXECUTE = "execute:async" + + +HEAVY_TYPES = frozenset({"image", "video", "audio", "mask"}) + + +def is_envelope(value: Any) -> bool: + return ( + isinstance(value, dict) + and isinstance(value.get("type"), str) + and value.get("type") in HEAVY_TYPES + and isinstance(value.get("encoding"), str) + ) + + +def decode_envelope_data(envelope: dict[str, Any]) -> bytes: + """Decode the inline base64 ``data`` of an envelope. + + URI envelopes are resolved out-of-band by the caller (see + ``proxy_node._deserialize_output``); by the time they reach this + function the bytes have already been fetched and re-stamped onto + ``data``. + """ + if "data" not in envelope: + raise ValueError( + "envelope has no inline 'data' — caller must resolve " + f"{sorted(k for k in envelope.keys() if k != 'data')!r} first" + ) + return base64.b64decode(envelope["data"]) + + +class RnpProtocolError(Exception): + """Raised when a server response can't be parsed as RNP/1, or when + a structured-error response is returned. Carries the parsed error + envelope so the caller can surface ``user_facing`` messages. + """ + + def __init__( + self, + message: str, + *, + code: str = ErrorCode.INTERNAL, + status: int | None = None, + user_facing: bool = False, + retryable: bool = False, + details: dict[str, Any] | None = None, + ) -> None: + super().__init__(message) + self.code = code + self.status = status + self.user_facing = user_facing + self.retryable = retryable + self.details = details + + @classmethod + def from_response_body( + cls, body: dict[str, Any] | None, status: int + ) -> "RnpProtocolError": + err = (body or {}).get("error") if isinstance(body, dict) else None + if not isinstance(err, dict): + return cls( + f"Unstructured error response (HTTP {status})", + code=ErrorCode.INTERNAL, + status=status, + ) + return cls( + err.get("message") or f"HTTP {status}", + code=err.get("code") or ErrorCode.INTERNAL, + status=status, + user_facing=bool(err.get("user_facing")), + retryable=bool(err.get("retryable")), + details=err.get("details") if isinstance(err.get("details"), dict) else None, + ) diff --git a/proxy_node.py b/proxy_node.py new file mode 100644 index 0000000..3e66606 --- /dev/null +++ b/proxy_node.py @@ -0,0 +1,1179 @@ +"""Build dynamic V3 ``IO.ComfyNode`` subclasses from RNP NodeDescriptors. + +The descriptor wire format is V3 ``Schema.get_v1_info()`` verbatim, so +per-input dicts are handed straight to V3's IO classes without +reinterpretation. ``_parse_input_spec`` knows STRING / INT / FLOAT / +BOOLEAN / COMBO / IMAGE / VIDEO / AUDIO / MASK; an input outside that +set causes the whole node to be skipped with a warning rather than +half-built. + +Hidden inputs (``auth_token_comfy_org`` / ``api_key_comfy_org`` / +``unique_id`` / ``prompt`` / ``extra_pnginfo`` / ``dynprompt``) are +declared on the schema so the executor populates ``cls.hidden``; auth +credentials are then lifted into outbound HTTP headers and everything +else flows into the request ``context`` map. + +Only ``execute`` is wired today; the validate / fingerprint / +check_lazy_status RNP methods fall through to ComfyUI's defaults until +their server endpoints exist. +""" +from __future__ import annotations + +import asyncio +import contextlib +import logging +import time +import uuid +from typing import Any + +from comfy_api.latest import IO +from server import PromptServer + +from comfy_api_nodes.util.client import ApiEndpoint, poll_op_raw +from comfy_api_nodes.util.common_exceptions import ProcessingInterrupted + +from . import client as rnp_client +from . import serialization +from .protocol import ( + DEFAULT_CANCEL_TIMEOUT_S, + DEFAULT_MAX_POLL_ATTEMPTS, + DEFAULT_MAX_RETRIES_PER_POLL, + DEFAULT_POLL_INTERVAL_S, + DEFAULT_RETRY_BACKOFF_PER_POLL, + DEFAULT_RETRY_DELAY_PER_POLL_S, + DEFAULT_TIMEOUT_PER_POLL_S, + ErrorCode, + ExecutionMode, + Header, + PROTOCOL_VERSION, + RnpProtocolError, + TaskStatus, + is_envelope, +) + + +# ComfyUI's frontend matches these strings literally on +# ``exception_message.includes(...)`` to fire dedicated dialogs, so the +# wording here must stay byte-identical with +# ``comfy_api_nodes.util.client._friendly_http_message``. +_AUTH_FAILED_MESSAGE = "Unauthorized: Please login first to use this node." +_INSUFFICIENT_CREDITS_MESSAGE = ( + "Payment Required: Please add credits to your account to use this node." +) +_ACCOUNT_PROBLEM_MESSAGE = ( + "There is a problem with your account. Please contact support@comfy.org." +) +_RATE_LIMIT_AFTER_RETRY_MESSAGE = ( + "Rate Limit Exceeded: The server returned 429 after all retry attempts. " + "Please wait and try again." +) +_PROVIDER_UNAVAILABLE_MESSAGE = ( + "The remote node provider is temporarily unavailable. Please try again " + "in a few minutes." +) +_SERVER_BUSY_MESSAGE = ( + "The remote node server is at capacity. Please try again in a moment." +) +_MAINTENANCE_MESSAGE = ( + "The remote node server is in maintenance. Please try again shortly." +) +_TASK_LOST_MESSAGE = ( + "The remote node server lost the running task and the resume attempt " + "did not succeed. Please re-run the workflow." +) + + +def _to_friendly_message(e: RnpProtocolError, node_id: str) -> str: + upstream = (e.details or {}).get("upstream_status") if e.details else None + if e.code == ErrorCode.AUTH_FAILED or upstream in (401, 403): + return _AUTH_FAILED_MESSAGE + if e.code == ErrorCode.INSUFFICIENT_CREDITS or upstream == 402: + return _INSUFFICIENT_CREDITS_MESSAGE + # 409 from the upstream gateway means "account in bad standing". + # comfy_api_nodes.util.client._friendly_http_message uses a dedicated + # support-contact prompt for this; mirror it byte-identically so the + # same frontend dialog fires. + if upstream == 409: + return _ACCOUNT_PROBLEM_MESSAGE + # 429 reaches the client only after the server has exhausted its own + # rate-limit retries — surface the post-retry copy so users don't + # think a single transient 429 leaked through. + if e.code == ErrorCode.RATE_LIMITED or upstream == 429: + return _RATE_LIMIT_AFTER_RETRY_MESSAGE + if e.code == ErrorCode.SERVER_BUSY: + return _SERVER_BUSY_MESSAGE + if e.code == ErrorCode.MAINTENANCE: + # Server-friendly message often embeds an "until" timestamp in + # the wire ``message`` — prefer it over the canned copy. + return str(e) if e.user_facing else _MAINTENANCE_MESSAGE + if e.code == ErrorCode.TASK_LOST: + return _TASK_LOST_MESSAGE + # PROVIDER_UNAVAILABLE with a retryable hint or 5xx upstream + # collapses to the generic "try again" copy. + if e.code == ErrorCode.PROVIDER_UNAVAILABLE and ( + upstream in (502, 503, 504) or e.retryable + ): + return _PROVIDER_UNAVAILABLE_MESSAGE + if e.user_facing: + return str(e) + return f"Remote node {node_id} failed ({e.code})" + + +def _coerce_pos(value: Any, default: float) -> float: + """Best-effort numeric coerce — falls back to ``default`` for anything + unparseable or non-positive (the descriptor came over the wire so we + don't trust it).""" + try: + f = float(value) + except (TypeError, ValueError): + return float(default) + return f if f > 0 else float(default) + + +def _coerce_pos_int(value: Any, default: int) -> int: + try: + i = int(value) + except (TypeError, ValueError): + return int(default) + return i if i > 0 else int(default) + + +def _extract_execution_policy(execution: dict[str, Any]) -> dict[str, Any]: + """Lift ``descriptor.remote.execution`` policy fields into a flat dict + of ``poll_op_raw`` kwargs. + + Defaults mirror ``comfy_api_nodes.util.client.poll_op_raw`` so an RNP + node with an empty ``execution`` block behaves identically to a + direct upstream call. The schema is the wire contract from + ``comfy_rnp_protocol.constants``: ``poll_interval_s``, + ``max_poll_attempts``, ``timeout_per_poll_s``, ``cancel_timeout_s``, + ``max_retries_per_poll``, ``retry_delay_per_poll_s``, + ``retry_backoff_per_poll``, ``retry: {max, delay_s, backoff, + retry_on}``, ``idempotency``. + """ + retry_block = execution.get("retry") if isinstance(execution.get("retry"), dict) else {} + return { + "poll_interval_s": _coerce_pos( + execution.get("poll_interval_s"), DEFAULT_POLL_INTERVAL_S, + ), + "max_poll_attempts": _coerce_pos_int( + execution.get("max_poll_attempts"), DEFAULT_MAX_POLL_ATTEMPTS, + ), + "timeout_per_poll_s": _coerce_pos( + execution.get("timeout_per_poll_s"), DEFAULT_TIMEOUT_PER_POLL_S, + ), + "cancel_timeout_s": _coerce_pos( + execution.get("cancel_timeout_s"), DEFAULT_CANCEL_TIMEOUT_S, + ), + "max_retries_per_poll": _coerce_pos_int( + execution.get("max_retries_per_poll"), DEFAULT_MAX_RETRIES_PER_POLL, + ), + "retry_delay_per_poll_s": _coerce_pos( + execution.get("retry_delay_per_poll_s"), DEFAULT_RETRY_DELAY_PER_POLL_S, + ), + "retry_backoff_per_poll": _coerce_pos( + execution.get("retry_backoff_per_poll"), DEFAULT_RETRY_BACKOFF_PER_POLL, + ), + # Submit-side retry block — currently unused by execute_remote / + # execute_remote_async (they hard-code max_retries=2 in + # rnp_client) but lifted out here so future steps can plumb it + # through without re-parsing the descriptor. + "submit_retry": { + "max": _coerce_pos_int(retry_block.get("max"), 0) + if retry_block.get("max") is not None else 0, + "delay_s": _coerce_pos(retry_block.get("delay_s"), 1.0), + "backoff": _coerce_pos(retry_block.get("backoff"), 2.0), + "retry_on": list(retry_block.get("retry_on") or []), + }, + "idempotency": execution.get("idempotency") or "none", + } + + +def _node_id_for_ticker(cls: type[IO.ComfyNode]) -> str | None: + """Lift ``cls.hidden.unique_id`` for the per-second progress ticker. + + Returns ``None`` if the hidden holder isn't populated (test paths + that build a class but don't run it through the executor). + """ + hidden = getattr(cls, "hidden", None) + if hidden is None: + return None + uid = getattr(hidden, "unique_id", None) + return str(uid) if uid is not None else None + + +async def _progress_ticker( + node_id: str | None, + estimated_duration_s: int | float | None, + *, + status_override: dict[str, Any] | None = None, +) -> None: + """Emit one ``send_progress_text`` per second until cancelled. + + Format mirrors ``comfy_api_nodes.util.client._display_time_progress`` + so the progress badge looks identical to local partner nodes: + ``Status: Processing\\nTime elapsed: s`` (with + ``(~Ns remaining)`` when the descriptor declared an estimate). + + ``status_override`` is a shared mutable holder + ``{"label": str|None, "expires_at": float|None}`` the caller can + mutate to swap "Processing" for a transient message like + "Server busy, retrying in 30s..." while a backpressure 503 is + being waited out. The ticker auto-clears the override once + ``time.monotonic() >= expires_at`` so a stale label doesn't linger + if the post-wait request takes longer than the announced + Retry-After window. + """ + if node_id is None: + return + started = time.monotonic() + while True: + elapsed = int(time.monotonic() - started) + if estimated_duration_s is not None and estimated_duration_s > 0: + remaining = max(0, int(estimated_duration_s) - elapsed) + time_line = f"Time elapsed: {elapsed}s (~{remaining}s remaining)" + else: + time_line = f"Time elapsed: {elapsed}s" + label = "Processing" + if status_override is not None: + expires_at = status_override.get("expires_at") + if expires_at is not None and time.monotonic() >= expires_at: + status_override["label"] = None + status_override["expires_at"] = None + override = status_override.get("label") + if override: + label = override + text = f"Status: {label}\n{time_line}" + with contextlib.suppress(Exception): + instance = PromptServer.instance + if instance is not None: + instance.send_progress_text(text, node_id) + try: + await asyncio.sleep(1.0) + except asyncio.CancelledError: + return + +log = logging.getLogger("comfy_remote_nodes.proxy_node") + + +def build_node_class( + node_id: str, + descriptor: dict[str, Any], + server_url: str, + *, + max_inline_bytes: int | None = None, +) -> type[IO.ComfyNode] | None: + """Build an ``IO.ComfyNode`` subclass for one descriptor. + + Returns ``None`` if the descriptor can't be parsed (e.g. an input + type ``_parse_input_spec`` doesn't recognise). The caller logs and + skips the node so one bad descriptor doesn't break the whole + extension. + """ + remote = descriptor.get("remote") or {} + if not isinstance(remote, dict): + log.warning("Descriptor %s missing 'remote' block — skipping", node_id) + return None + endpoints = remote.get("endpoints") or {} + execute_endpoint = endpoints.get("execute") + if not isinstance(execute_endpoint, dict) or not execute_endpoint.get("path"): + log.warning("Descriptor %s has no remote.endpoints.execute — skipping", node_id) + return None + + inputs, hidden_names = _parse_inputs(descriptor) + if inputs is None: + return None + outputs = _parse_outputs(descriptor) + if outputs is None: + return None + + schema_hash = remote.get("schema_hash") + execute_path = execute_endpoint["path"] + execution = remote.get("execution") or {} + timeout_s = float( + execution.get("hard_timeout_s") + or execution.get("soft_timeout_s") + or 600.0 + ) + estimated_duration_s = execution.get("estimated_duration_s") + if not isinstance(estimated_duration_s, (int, float)): + estimated_duration_s = None + execution_mode = execution.get("mode") or ExecutionMode.REQUEST_RESPONSE + execution_policy = _extract_execution_policy(execution) + # ``descriptor.remote.url_fetch`` declares per-output-type fetch + # posture for envelopes that arrive with ``uri`` instead of inline + # ``data``. Empty / absent → all outputs must inline their bytes. + url_fetch_policy = ( + remote.get("url_fetch") if isinstance(remote.get("url_fetch"), dict) else {} + ) + + display_name = descriptor.get("display_name") or node_id + category = descriptor.get("category") or "remote" + description = descriptor.get("description") or "" + is_api_node = bool(descriptor.get("api_node", False)) + is_experimental = bool(descriptor.get("experimental", False)) + is_deprecated = bool(descriptor.get("deprecated", False)) + + schema_inputs = list(inputs) + schema_outputs = list(outputs) + hidden_decls = _hidden_decls_from_names(hidden_names, is_api_node) + + # Forward every other V3 NodeInfoV1 field the descriptor surfaces. + # Anything not constructed here is silently dropped; the + # ``_parse_input_spec`` / ``_parse_outputs`` warnings are the + # signal a deployment needs to add a missing type to this builder. + is_input_list = bool(descriptor.get("is_input_list", False)) + is_output_node = bool(descriptor.get("output_node", False)) + is_dev_only = bool(descriptor.get("dev_only", False)) + has_intermediate_output = bool(descriptor.get("has_intermediate_output", False)) + search_aliases = descriptor.get("search_aliases") or [] + essentials_category = descriptor.get("essentials_category") + price_badge = _build_price_badge(descriptor.get("price_badge")) + + class RemoteProxyNode(IO.ComfyNode): + @classmethod + def define_schema(cls) -> IO.Schema: + return IO.Schema( + node_id=node_id, + display_name=display_name, + category=category, + description=description, + inputs=schema_inputs, + outputs=schema_outputs, + hidden=hidden_decls, + is_input_list=is_input_list, + is_output_node=is_output_node, + is_api_node=is_api_node, + is_experimental=is_experimental, + is_deprecated=is_deprecated, + is_dev_only=is_dev_only, + has_intermediate_output=has_intermediate_output, + search_aliases=list(search_aliases), + essentials_category=essentials_category, + price_badge=price_badge, + ) + + @classmethod + async def execute(cls, **kwargs: Any) -> IO.NodeOutput: + return await _execute_remote( + cls, + node_id=node_id, + server_url=server_url, + execute_path=execute_path, + schema_hash=schema_hash, + timeout_s=timeout_s, + outputs_meta=schema_outputs, + hidden_names=hidden_names, + inputs=kwargs, + max_inline_bytes=max_inline_bytes, + estimated_duration_s=estimated_duration_s, + execution_mode=execution_mode, + execution_policy=execution_policy, + url_fetch_policy=url_fetch_policy, + ) + + RemoteProxyNode.__name__ = node_id + RemoteProxyNode.__qualname__ = node_id + RemoteProxyNode.RELATIVE_PYTHON_MODULE = "custom_nodes.comfy_remote_nodes" + return RemoteProxyNode + + +# --------------------------------------------------------------------------- +# Input parsing +# --------------------------------------------------------------------------- + +# The hidden input markers V3 emits in get_v1_info() — clients forward +# their values through the request ``context`` instead of treating them +# as user-facing inputs. +_HIDDEN_MARKERS = { + "AUTH_TOKEN_COMFY_ORG": IO.Hidden.auth_token_comfy_org, + "API_KEY_COMFY_ORG": IO.Hidden.api_key_comfy_org, + "UNIQUE_ID": IO.Hidden.unique_id, + "PROMPT": IO.Hidden.prompt, + "EXTRA_PNGINFO": IO.Hidden.extra_pnginfo, + "DYNPROMPT": IO.Hidden.dynprompt, +} + + +def _parse_inputs(descriptor: dict[str, Any]) -> tuple[list[Any] | None, list[str]]: + """Return ``(input_objs, hidden_input_names)``. + + Returns ``(None, [])`` if any required/optional input can't be + parsed (caller logs and skips). + """ + input_def = descriptor.get("input") or {} + input_order = descriptor.get("input_order") or {} + out: list[Any] = [] + for kind, optional in (("required", False), ("optional", True)): + order = input_order.get(kind) or list((input_def.get(kind) or {}).keys()) + defs = input_def.get(kind) or {} + for name in order: + spec = defs.get(name) + if spec is None: + continue + inp = _parse_input_spec(name, spec, optional=optional) + if inp is None: + log.warning( + "Skipping descriptor: unsupported input %r (%s)", + name, spec, + ) + return None, [] + out.append(inp) + hidden_names = list((input_def.get("hidden") or {}).keys()) + return out, hidden_names + + +def _parse_input_spec(name: str, spec: list[Any], optional: bool) -> Any | None: + """Map one V3-shaped ``[io_type, options_dict]`` entry to a V3 IO Input.""" + if not isinstance(spec, (list, tuple)) or len(spec) < 1: + return None + io_type = spec[0] + options: dict[str, Any] = spec[1] if len(spec) > 1 and isinstance(spec[1], dict) else {} + common = { + "tooltip": options.get("tooltip"), + "optional": optional, + } + if options.get("advanced") is not None: + common["advanced"] = bool(options.get("advanced")) + + if isinstance(io_type, list): + # Legacy V1-style combo: io_type *is* the options list. + return IO.Combo.Input(name, options=list(io_type), + default=options.get("default"), **common) + if io_type == "COMBO": + opts = options.get("options") + if not isinstance(opts, (list, tuple)): + return None + return IO.Combo.Input(name, options=list(opts), + default=options.get("default"), **common) + if io_type == "STRING": + return IO.String.Input( + name, + default=options.get("default", ""), + multiline=bool(options.get("multiline", False)), + **common, + ) + if io_type == "INT": + return IO.Int.Input( + name, + default=options.get("default", 0), + min=options.get("min", 0), + max=options.get("max", 2147483647), + step=options.get("step", 1), + **common, + ) + if io_type == "FLOAT": + return IO.Float.Input( + name, + default=options.get("default", 0.0), + min=options.get("min", 0.0), + max=options.get("max", 1.0), + step=options.get("step", 0.01), + **common, + ) + if io_type == "BOOLEAN": + return IO.Boolean.Input( + name, default=bool(options.get("default", False)), **common, + ) + if io_type == "IMAGE": + return IO.Image.Input(name, **common) + if io_type == "VIDEO": + return IO.Video.Input(name, **common) + if io_type == "AUDIO": + return IO.Audio.Input(name, **common) + if io_type == "MASK": + return IO.Mask.Input(name, **common) + return None + + +def _build_price_badge(wire: Any) -> Any | None: + """Convert the wire-format ``price_badge`` dict back to an + ``IO.PriceBadge`` instance. + + The wire shape (from ``PriceBadge.as_dict()``) carries + ``depends_on.widgets`` as a list of ``{name, type}`` objects so the + frontend can resolve types; the V3 constructor only needs the + ``name`` strings since it re-derives types from the schema's actual + inputs at ``as_dict()`` time. Inputs / input_groups round-trip as-is. + + Returns ``None`` if the descriptor has no ``price_badge`` field. + """ + if not isinstance(wire, dict): + return None + expr = wire.get("expr") + if not isinstance(expr, str) or not expr.strip(): + return None + dep = wire.get("depends_on") or {} + widgets = [ + w["name"] for w in (dep.get("widgets") or []) if isinstance(w, dict) and "name" in w + ] + inputs = list(dep.get("inputs") or []) + input_groups = list(dep.get("input_groups") or []) + engine = wire.get("engine", "jsonata") + try: + return IO.PriceBadge( + expr=expr, + depends_on=IO.PriceBadgeDepends( + widgets=widgets, inputs=inputs, input_groups=input_groups, + ), + engine=engine, + ) + except Exception as e: # noqa: BLE001 + log.warning("Skipping price_badge — could not construct: %s", e) + return None + + +def _hidden_decls_from_names(hidden_names: list[str], is_api_node: bool) -> list[Any]: + """Build the ``hidden=`` list for IO.Schema. + + The descriptor lists each hidden input by name (e.g. + ``auth_token_comfy_org``); the ``_HIDDEN_MARKERS`` table above + resolves those names back to V3 ``IO.Hidden`` enum members. Names + that aren't in the table are dropped with a debug log so a server + that adds a new hidden marker doesn't break older clients. + """ + out: list[Any] = [] + seen: set[Any] = set() + for name in hidden_names: + marker = getattr(IO.Hidden, name, None) + if marker is None: + log.debug("Unknown hidden input %s — dropping", name) + continue + if marker not in seen: + out.append(marker) + seen.add(marker) + # IO.Schema.finalize() will add auth/api hidden if is_api_node, so + # we don't need to add them again here. + if not is_api_node: + # Non-api-node descriptors that happened to advertise auth + # hidden are forwarded as-is. + pass + return out + + +# --------------------------------------------------------------------------- +# Output parsing +# --------------------------------------------------------------------------- + +def _parse_outputs(descriptor: dict[str, Any]) -> list[Any] | None: + output_types = descriptor.get("output") or [] + output_names = descriptor.get("output_name") or output_types + output_is_list = descriptor.get("output_is_list") or [False] * len(output_types) + output_tooltips = descriptor.get("output_tooltips") or [None] * len(output_types) + out: list[Any] = [] + for i, io_type in enumerate(output_types): + name = output_names[i] if i < len(output_names) else io_type + is_list = bool(output_is_list[i]) if i < len(output_is_list) else False + tip = output_tooltips[i] if i < len(output_tooltips) else None + cls = _OUTPUT_CLASSES.get(io_type) + if cls is None: + log.warning( + "Skipping descriptor: unsupported output type %r", io_type, + ) + return None + out.append(cls(display_name=name, is_output_list=is_list, tooltip=tip)) + return out + + +_OUTPUT_CLASSES = { + "VIDEO": IO.Video.Output, + "IMAGE": IO.Image.Output, + "AUDIO": IO.Audio.Output, + "MASK": IO.Mask.Output, + "STRING": IO.String.Output, + "INT": IO.Int.Output, + "FLOAT": IO.Float.Output, + "BOOLEAN": IO.Boolean.Output, +} + + +# --------------------------------------------------------------------------- +# Execute dispatch +# --------------------------------------------------------------------------- + +def _extract_progress_pct(response: dict[str, Any]) -> int | None: + """Map the RNP ``progress`` block to the 0-100 int that + ``poll_op_raw`` feeds to its (single) ``ProgressBar(100)``. + + Returns ``None`` when the server doesn't yet have progress to + report so the bar stays at its current position rather than + snapping back to 0. + """ + progress = response.get("progress") + if not isinstance(progress, dict): + return None + value = progress.get("value") + total = progress.get("max") + if not isinstance(value, (int, float)) or not isinstance(total, (int, float)) or total <= 0: + return None + pct = int(round((float(value) / float(total)) * 100)) + return max(0, min(100, pct)) + + +def _rnp_protocol_headers(extra: dict[str, str] | None = None) -> dict[str, str]: + """Headers every RNP request carries — protocol-version + client + version. Pass ``extra`` to merge per-request headers (e.g. the + Idempotency-Key the server uses to authorize a TASK_LOST resume, + or the freshly-rotated auth token a long-running poll wants to + forward to the server).""" + headers = { + Header.PROTOCOL_VERSION: PROTOCOL_VERSION, + Header.CLIENT_VERSION: rnp_client.CLIENT_VERSION, + } + if extra: + headers.update(extra) + return headers + + +async def _run_async_polling( + cls: type[IO.ComfyNode], + *, + node_id: str, + server_url: str, + schema_hash: str | None, + payload_inputs: dict[str, Any], + context: dict[str, Any], + auth_headers: dict[str, str], + idempotency_key: str, + poll_interval_s: float, + estimated_duration_s: int | float | None, + max_poll_attempts: int, + timeout_per_poll_s: float, + cancel_timeout_s: float, + max_retries_per_poll: int, + retry_delay_per_poll_s: float, + retry_backoff_per_poll: float, +) -> dict[str, Any]: + """Drive the ``execute_async`` + ``poll`` flow via ``poll_op_raw``. + + Submits the task with ``execute_async``, then hands the polling + loop off to ``comfy_api_nodes.util.client.poll_op_raw`` so RNP + inherits the same per-poll retry/backoff, queued-vs-active + accounting, max-poll-attempts cap, ``ProgressBar`` plumbing, + ticker UI, and cancel-endpoint propagation as local partner + nodes. All of DONE/ERROR/CANCELLED are mapped to poll_op_raw's + "completed" set so the function returns the terminal poll-response + dict and we dispatch on ``status`` here — that preserves + ``RnpProtocolError`` typing for ERROR and ``ProcessingInterrupted`` + for CANCELLED instead of flattening both to + ``Exception("Task failed: …")``. + + On ``TASK_LOST`` (the server lost the task store between submit + and poll, typically a restart), the function resubmits exactly + once with the same idempotency key and resumes polling. A second + ``TASK_LOST`` from the resume gives up — at that point the failure + is structural rather than transient. + """ + resume_attempted = False + while True: + submit = await rnp_client.execute_remote_async( + server_url, + schema_hash, + node_id, + payload_inputs, + context, + extra_headers=auth_headers or None, + idempotency_key=idempotency_key, + node_cls=cls, + ) + task_id = submit["task_id"] + interval = float(submit.get("poll_interval") or poll_interval_s) + try: + response = await _poll_until_terminal( + cls, + node_id=node_id, + server_url=server_url, + task_id=task_id, + auth_headers=auth_headers, + idempotency_key=idempotency_key, + poll_interval_s=interval, + estimated_duration_s=estimated_duration_s, + max_poll_attempts=max_poll_attempts, + timeout_per_poll_s=timeout_per_poll_s, + cancel_timeout_s=cancel_timeout_s, + max_retries_per_poll=max_retries_per_poll, + retry_delay_per_poll_s=retry_delay_per_poll_s, + retry_backoff_per_poll=retry_backoff_per_poll, + ) + except RnpProtocolError as e: + if e.code == ErrorCode.TASK_LOST and not resume_attempted: + resume_attempted = True + _send_progress_text( + cls, + "Server restarted, resuming…", + ) + log.warning( + "RNP %s: TASK_LOST during poll (task_id=%s); " + "resubmitting with idempotency_key=%s", + node_id, task_id, idempotency_key, + ) + continue + raise + return _interpret_terminal_response(task_id, response) + + +def _send_progress_text(cls: type[IO.ComfyNode], text: str) -> None: + """Best-effort progress badge update — invisible when the executor + isn't running this class through PromptServer (e.g. unit tests).""" + node_id = _node_id_for_ticker(cls) + if node_id is None: + return + with contextlib.suppress(Exception): + instance = PromptServer.instance + if instance is not None: + instance.send_progress_text(text, node_id) + + +async def _poll_until_terminal( + cls: type[IO.ComfyNode], + *, + node_id: str, + server_url: str, + task_id: str, + auth_headers: dict[str, str], + idempotency_key: str, + poll_interval_s: float, + estimated_duration_s: int | float | None, + max_poll_attempts: int, + timeout_per_poll_s: float, + cancel_timeout_s: float, + max_retries_per_poll: int, + retry_delay_per_poll_s: float, + retry_backoff_per_poll: float, +) -> dict[str, Any]: + """Run ``poll_op_raw`` against an existing task_id and return the + terminal poll-response dict. Wraps the per-request header build so + the resume path in ``_run_async_polling`` can reuse it.""" + # Forward the idempotency key on every poll/cancel: the server + # uses it to gate TASK_LOST recovery (a typo'd task_id without an + # idempotency key falls through to a generic NOT_FOUND so it + # cannot trigger a phantom resubmit). Auth headers ride alongside + # so a token rotated mid-task picks up on the next upstream call + # the server makes. + per_request_headers = _rnp_protocol_headers({ + Header.IDEMPOTENCY_KEY: idempotency_key, + **(auth_headers or {}), + }) + poll_endpoint = ApiEndpoint( + f"rnp/v1/tasks/{task_id}", "GET", + headers=per_request_headers, + ) + cancel_endpoint = ApiEndpoint( + f"rnp/v1/tasks/{task_id}/cancel", "POST", + headers=per_request_headers, + ) + return await poll_op_raw( + cls, + poll_endpoint, + status_extractor=lambda r: r.get("status"), + progress_extractor=_extract_progress_pct, + completed_statuses=[TaskStatus.DONE, TaskStatus.ERROR, TaskStatus.CANCELLED], + failed_statuses=[], + queued_statuses=[TaskStatus.PENDING], + poll_interval=poll_interval_s, + max_poll_attempts=max_poll_attempts, + timeout_per_poll=timeout_per_poll_s, + max_retries_per_poll=max_retries_per_poll, + retry_delay_per_poll=retry_delay_per_poll_s, + retry_backoff_per_poll=retry_backoff_per_poll, + estimated_duration=int(estimated_duration_s) if estimated_duration_s else None, + cancel_endpoint=cancel_endpoint, + cancel_timeout=cancel_timeout_s, + base_url=server_url, + auth_headers={}, + error_parser=rnp_client._parse_rnp_error, + is_rate_limited=rnp_client._is_rnp_backpressure, + rate_limit_label=rnp_client._rnp_rate_limit_label, + ) + + +def _interpret_terminal_response( + task_id: str, response: dict[str, Any], +) -> dict[str, Any]: + """Map a terminal poll-response dict to the same ``{"outputs": [...]}`` + shape sync ``execute`` returns, or raise the appropriate typed + exception for ERROR / CANCELLED.""" + status = response.get("status") + if status == TaskStatus.DONE: + outputs = response.get("outputs") + if not isinstance(outputs, list): + raise RnpProtocolError( + f"task {task_id} done but no outputs list: {response!r}", + ) + return {"outputs": outputs} + if status == TaskStatus.ERROR: + err_body = response.get("exception") + if not isinstance(err_body, dict): + raise RnpProtocolError( + f"task {task_id} errored without exception body", + ) + raise RnpProtocolError.from_response_body(err_body, 500) + if status == TaskStatus.CANCELLED: + raise ProcessingInterrupted(f"task {task_id} cancelled") + raise RnpProtocolError( + f"task {task_id} returned unrecognised terminal status: {status!r}", + ) + + +async def _execute_remote( + cls: type[IO.ComfyNode], + *, + node_id: str, + server_url: str, + execute_path: str, + schema_hash: str | None, + timeout_s: float, + outputs_meta: list[Any], + hidden_names: list[str], + inputs: dict[str, Any], + max_inline_bytes: int | None = None, + estimated_duration_s: int | float | None = None, + execution_mode: str = ExecutionMode.REQUEST_RESPONSE, + execution_policy: dict[str, Any] | None = None, + url_fetch_policy: dict[str, Any] | None = None, +) -> IO.NodeOutput: + policy = execution_policy or _extract_execution_policy({}) + url_fetch_policy = url_fetch_policy or {} + # Strip hidden inputs from the user-supplied kwargs. Auth credentials + # (``auth_token_comfy_org`` / ``api_key_comfy_org``) travel as standard + # HTTP headers so the server can forward them to api.comfy.org without + # parsing the body; everything else rides in the request ``context``. + context: dict[str, Any] = {} + auth_headers: dict[str, str] = {} + hidden_holder = getattr(cls, "hidden", None) + for hname in hidden_names: + # Prefer cls.hidden (the standard executor path); fall back to a + # value piped directly into kwargs for advanced workflows that + # wire strings into hidden slots manually. Either way, strip the + # name from the user-input kwargs so it doesn't leak into the + # request body. + hidden_value = ( + getattr(hidden_holder, hname, None) if hidden_holder is not None else None + ) + piped_value = inputs.pop(hname, None) + value = hidden_value if hidden_value is not None else piped_value + if value is None: + continue + if hname == "auth_token_comfy_org": + auth_headers["Authorization"] = f"Bearer {value}" + elif hname == "api_key_comfy_org": + auth_headers["X-API-KEY"] = str(value) + else: + context[hname] = value + + # Encode heavy-typed inputs (IMAGE/MASK/AUDIO) as RNP value + # envelopes; scalars pass through. Oversized envelopes are + # uploaded out-of-band and rewritten to ``{..., uri: }``. + payload_inputs = await _encode_inputs( + inputs, + server_url=server_url, + max_inline_bytes=max_inline_bytes, + auth_headers=auth_headers, + ) + + ticker_node_id = _node_id_for_ticker(cls) + # Stable per-execute() idempotency key — when the underlying HTTP + # client retries the POST after a connection drop, the server sees + # the same key and returns the original task_id rather than + # spawning a duplicate run. + idempotency_key = uuid.uuid4().hex + + async def _run_with_ticker( + execute_schema_hash: str | None, + ) -> dict[str, Any]: + if execution_mode == ExecutionMode.ASYNC_POLLING: + # poll_op_raw runs its own _display_time_progress ticker, + # so we don't start a parallel one here — the sync path + # below still uses the local ticker since execute_remote + # blocks on a single HTTP request without UI feedback. + return await _run_async_polling( + cls, + node_id=node_id, + server_url=server_url, + schema_hash=execute_schema_hash, + payload_inputs=payload_inputs, + context=context, + auth_headers=auth_headers, + idempotency_key=idempotency_key, + poll_interval_s=policy["poll_interval_s"], + estimated_duration_s=estimated_duration_s, + max_poll_attempts=policy["max_poll_attempts"], + timeout_per_poll_s=policy["timeout_per_poll_s"], + cancel_timeout_s=policy["cancel_timeout_s"], + max_retries_per_poll=policy["max_retries_per_poll"], + retry_delay_per_poll_s=policy["retry_delay_per_poll_s"], + retry_backoff_per_poll=policy["retry_backoff_per_poll"], + ) + # Shared mutable holder the local ticker watches so a + # SERVER_BUSY/MAINTENANCE retry inside execute_remote can swap + # the per-second "Status: Processing" line for the friendly + # backpressure copy without restarting the ticker. The ticker + # auto-expires the override once ``expires_at`` passes so a + # stale "Server busy, retrying in 30s..." doesn't linger if the + # post-wait request itself takes longer than the announced + # Retry-After window. + status_override: dict[str, Any] = {"label": None, "expires_at": None} + + def _local_rate_limit_label( + status: int, body: Any, retry_after_s: float, + ) -> str | None: + label = rnp_client._rnp_rate_limit_label(status, body, retry_after_s) + status_override["label"] = label # may be None for unknown codes + status_override["expires_at"] = ( + time.monotonic() + max(1.0, retry_after_s) if label else None + ) + return label + + ticker = asyncio.create_task( + _progress_ticker( + ticker_node_id, + estimated_duration_s, + status_override=status_override, + ), + ) + try: + return await rnp_client.execute_remote( + server_url, + execute_schema_hash, + execute_path, + payload_inputs, + context, + timeout_s=timeout_s, + extra_headers=auth_headers or None, + node_cls=cls, + rate_limit_label=_local_rate_limit_label, + ) + finally: + ticker.cancel() + with contextlib.suppress(BaseException): + await ticker + + try: + body = await _run_with_ticker(schema_hash) + except RnpProtocolError as e: + # Schema-hash mismatch → re-fetch object_info, find the new + # hash for this node, and retry once. This recovers from the + # common deploy ordering of "server updated descriptor ⇒ + # client still using cached one" without forcing a ComfyUI + # restart. + if e.code == ErrorCode.SCHEMA_HASH_MISMATCH: + new_hash = await _refresh_schema_hash(server_url, node_id) + if new_hash and new_hash != schema_hash: + log.info( + "RNP %s: schema_hash changed (%s → %s); retrying execute", + node_id, schema_hash, new_hash, + ) + try: + body = await _run_with_ticker(new_hash) + except RnpProtocolError as e2: + raise RuntimeError( + _to_friendly_message(e2, node_id), + ) from e2 + else: + raise RuntimeError( + f"Remote node {node_id}: schema-hash mismatch and no " + f"updated descriptor found upstream" + ) from e + else: + log.warning("RNP execute %s failed: %s [%s]", node_id, e, e.code) + raise RuntimeError(_to_friendly_message(e, node_id)) from e + + outputs = body.get("outputs") + if not isinstance(outputs, list): + raise RuntimeError( + f"Remote node {node_id} returned no 'outputs' list" + ) + if len(outputs) != len(outputs_meta): + raise RuntimeError( + f"Remote node {node_id} returned {len(outputs)} outputs; " + f"schema declared {len(outputs_meta)}" + ) + + decoded = [] + for value, meta in zip(outputs, outputs_meta): + decoded.append(await _deserialize_output( + value, + meta, + url_fetch_policy=url_fetch_policy, + server_url=server_url, + auth_headers=auth_headers, + node_cls=cls, + )) + return IO.NodeOutput(*decoded) + + +async def _refresh_schema_hash(server_url: str, node_id: str) -> str | None: + """Re-fetch ``/object_info`` and return the freshly published + ``remote.schema_hash`` for ``node_id`` (or ``None`` if the node has + been removed upstream). + """ + try: + body, _etag, _max_age = await rnp_client.fetch_object_info(server_url) + except Exception as e: # noqa: BLE001 + log.warning("schema-hash refresh failed for %s: %s", node_id, e) + return None + if not isinstance(body, dict): + return None + descriptor = body.get(node_id) + if not isinstance(descriptor, dict): + return None + remote = descriptor.get("remote") or {} + return remote.get("schema_hash") if isinstance(remote, dict) else None + + +async def _deserialize_output( + value: Any, + meta: Any, + *, + url_fetch_policy: dict[str, Any], + server_url: str, + auth_headers: dict[str, str], + node_cls: type[IO.ComfyNode] | None, +) -> Any: + """Convert one output value (envelope or scalar) to its native Python form. + + Inline envelopes (with ``data``) decode immediately. URI envelopes + look up their fetch policy from ``url_fetch_policy`` (keyed by the + envelope's ``type``) and fetch the bytes via ``sync_op_raw`` so the + GET inherits ``ProcessingInterrupted`` cancellation, retry/backoff + on transient errors, and the per-policy timeout. Scalars pass + through unchanged.""" + if not is_envelope(value): + # Scalars pass through unchanged — the schema's output type is + # advisory here, the value already matches the native ComfyUI + # type for STRING / INT / FLOAT / BOOLEAN. + return value + if "data" not in value and "uri" in value: + env_type = value.get("type", "") + policy = url_fetch_policy.get(env_type) + if not isinstance(policy, dict): + raise RnpProtocolError( + f"Server returned a URL envelope for output type {env_type!r} " + "but the descriptor declared no remote.url_fetch policy " + "for that type", + code=ErrorCode.INTERNAL, + ) + raw = await _fetch_url_envelope( + value, policy, + auth_headers=auth_headers, + node_cls=node_cls, + ) + # Re-stamp the envelope with inline ``data`` so the existing + # decoder dispatch in ``serialization`` sees a homogeneous + # input shape regardless of how the bytes arrived. + import base64 as _b64 + value = {k: v for k, v in value.items() if k != "uri"} + value["data"] = _b64.b64encode(raw).decode("ascii") + return serialization.decode_envelope(value) + + +async def _fetch_url_envelope( + envelope: dict[str, Any], + policy: dict[str, Any], + *, + auth_headers: dict[str, str], + node_cls: type[IO.ComfyNode] | None, +) -> bytes: + """GET ``envelope["uri"]`` via ``sync_op_raw`` honoring the descriptor's + url_fetch policy (retries, backoff, timeout, optional auth). + + ``url_kind`` drives the auth posture: ``presigned`` / ``public`` → + no auth header (the URL embeds its own credentials or needs none); + ``comfy_api`` → forward the same Authorization / X-API-KEY the + polls use. ``auth_required`` overrides ``url_kind`` when explicitly + set to ``True``. + """ + from comfy_api_nodes.util.client import ApiEndpoint, sync_op_raw + + uri = envelope["uri"] + url_kind = policy.get("url_kind") or "presigned" + auth_required = bool(policy.get( + "auth_required", + url_kind == "comfy_api", + )) + fetch_auth: dict[str, str] = ( + dict(auth_headers) if (auth_required and auth_headers) else {} + ) + timeout_s = float(policy.get("timeout_s") or 300) + max_retries = int(policy.get("max_retries") or 3) + retry_delay = float(policy.get("retry_delay_s") or 1.0) + retry_backoff = float(policy.get("retry_backoff") or 2.0) + + raw = await sync_op_raw( + node_cls, + ApiEndpoint(uri, "GET"), + timeout=timeout_s, + max_retries=max_retries, + retry_delay=retry_delay, + retry_backoff=retry_backoff, + wait_label="Downloading", + as_binary=True, + final_label_on_success=None, + monitor_progress=False, + auth_headers=fetch_auth, + ) + if not isinstance(raw, (bytes, bytearray)): + raise RnpProtocolError( + f"url_fetch GET {uri!r} returned non-bytes payload", + code=ErrorCode.INTERNAL, + ) + return bytes(raw) + + +async def _encode_inputs( + inputs: dict[str, Any], + *, + server_url: str | None = None, + max_inline_bytes: int | None = None, + auth_headers: dict[str, str] | None = None, +) -> dict[str, Any]: + """Convert in-place: replace heavy-typed input values with the + matching RNP value envelope. Scalars pass through unchanged. + + The duck-typing rules: + + * ``torch.Tensor`` rank 4 (B,H,W,C) → image envelope. + * ``torch.Tensor`` rank 3 with last-dim != 3/4 → mask envelope (B,H,W). + Rank-2 (H,W) tensors are also treated as masks. + * ``dict`` with ``waveform`` + ``sample_rate`` → audio envelope. + * Everything else passes through as-is (already JSON-serializable). + + When ``max_inline_bytes`` is set, any envelope whose inline payload + exceeds the threshold is externalized via ``POST /rnp/v1/uploads`` + (see ``serialization.maybe_externalize``). + """ + out: dict[str, Any] = {} + for name, value in inputs.items(): + encoded = _encode_one(name, value) + # Only envelope-shaped values can be externalized; scalars pass + # straight through ``maybe_externalize`` unchanged. + if serialization.is_envelope(encoded): + encoded = await serialization.maybe_externalize( + encoded, + server_url=server_url, + max_inline_bytes=max_inline_bytes, + auth_headers=auth_headers, + ) + out[name] = encoded + return out + + +def _encode_one(name: str, value: Any) -> Any: + if serialization.is_audio_input(value): + return serialization.encode_audio_input(value) + if not serialization._is_torch_tensor(value): + return value + rank = value.dim() + last_dim = int(value.shape[-1]) if rank >= 1 else 0 + if rank == 4 or (rank == 3 and last_dim in (3, 4)): + # Per-image tensors: a length-1 batch is the common case; longer + # batches send only the first frame (the framework hasn't asked + # for batch-of-envelopes semantics yet). + return serialization.encode_image_tensor(value) + if rank in (2, 3): + return serialization.encode_mask_tensor(value) + log.warning( + "Input %s: unrecognised tensor rank %d — sending as-is", + name, rank, + ) + return value diff --git a/registry.py b/registry.py new file mode 100644 index 0000000..8818ba8 --- /dev/null +++ b/registry.py @@ -0,0 +1,177 @@ +"""ComfyExtension entry point — handshake with the RNP server and +register dynamic node classes with ComfyUI. + +Reads the server URL from ``RNP_SERVER_URL``. When the env var is unset +the extension registers zero nodes and logs a single "no server +configured" line — never crashes ComfyUI startup. + +If the server is unreachable or the manifest/object_info fetch fails, +the extension degrades silently with a log line so a flaky deployment +doesn't take ComfyUI down with it. +""" +from __future__ import annotations + +import logging +import os +from typing import Any + +from comfy_api.latest import ComfyExtension, IO +from typing_extensions import override + +from . import client as rnp_client +from .protocol import PROTOCOL_VERSION, RnpProtocolError +from .proxy_node import build_node_class + +log = logging.getLogger("comfy_remote_nodes") + + +SERVER_URL_ENV = "RNP_SERVER_URL" + + +class _RemoteNodesExtension(ComfyExtension): + """Concrete ComfyExtension that pulls node descriptors from one + RNP server URL and exposes the resulting dynamic classes to + ComfyUI.""" + + def __init__(self, server_url: str) -> None: + self._server_url = server_url + self._node_classes: list[type[IO.ComfyNode]] = [] + # ETag + max-age from the latest object_info fetch — used by a + # future re-discovery loop to do conditional GETs. + self._object_info_etag: str | None = None + self._object_info_max_age: int | None = None + # Inline-payload cap from the manifest. Populated in on_load(). + self._max_inline_bytes: int | None = None + + @override + async def on_load(self) -> None: + try: + manifest = await rnp_client.fetch_manifest(self._server_url) + except RnpProtocolError as e: + log.warning( + "RNP manifest fetch from %s failed (%s): %s", + self._server_url, e.code, e, + ) + return + except Exception as e: # noqa: BLE001 + log.warning( + "RNP manifest fetch from %s failed: %s", + self._server_url, e, + ) + return + + server_proto = manifest.get("protocol_version") or "" + if not _major_compatible(server_proto, PROTOCOL_VERSION): + log.warning( + "RNP server %s advertises protocol %s; client speaks %s — skipping", + self._server_url, server_proto, PROTOCOL_VERSION, + ) + return + # Cap on inline base64 payloads in /execute requests. Envelopes + # larger than this are uploaded via /rnp/v1/uploads and the + # payload is rewritten to a download URL. ``None`` disables. + self._max_inline_bytes = manifest.get("max_inline_payload_bytes") + log.info( + "RNP server %s ok: provider=%s protocol=%s max_inline_bytes=%s", + self._server_url, + (manifest.get("provider") or {}).get("id"), + server_proto, + self._max_inline_bytes, + ) + + try: + object_info, etag, max_age = await rnp_client.fetch_object_info( + self._server_url + ) + except RnpProtocolError as e: + log.warning( + "RNP object_info fetch from %s failed (%s): %s", + self._server_url, e.code, e, + ) + return + except Exception as e: # noqa: BLE001 + log.warning( + "RNP object_info fetch from %s failed: %s", + self._server_url, e, + ) + return + + # First-load case: 304 should not happen because we sent no + # If-None-Match. Defensive check anyway. + if object_info is None: + log.warning( + "RNP object_info from %s returned 304 with no cached body", + self._server_url, + ) + return + + # Stash the cache state on the extension instance so a future + # rediscovery loop can issue cheap conditional GETs against + # ``/object_info``. + self._object_info_etag = etag + self._object_info_max_age = max_age + log.debug( + "RNP object_info etag=%s max_age=%s", + etag, max_age, + ) + + for node_id, descriptor in object_info.items(): + if not isinstance(descriptor, dict): + log.warning("Skipping %s: descriptor is not a dict", node_id) + continue + try: + node_cls = build_node_class( + node_id, + descriptor, + self._server_url, + max_inline_bytes=self._max_inline_bytes, + ) + except Exception as e: # noqa: BLE001 + log.warning("Failed to build remote node %s: %s", node_id, e) + continue + if node_cls is not None: + self._node_classes.append(node_cls) + log.info("Registered remote node: %s", node_id) + log.info( + "comfy_remote_nodes ready: %d node(s) loaded from %s", + len(self._node_classes), self._server_url, + ) + + @override + async def get_node_list(self) -> list[type[IO.ComfyNode]]: + return list(self._node_classes) + + +def _major_compatible(server_version: str, client_version: str) -> bool: + """True iff the major component of two semver strings matches. + + A major-version mismatch is a hard reject; minor / patch + differences are tolerated (additive evolution only). + """ + if not server_version or not client_version: + return False + try: + return server_version.split(".")[0] == client_version.split(".")[0] + except Exception: # noqa: BLE001 + return False + + +async def comfy_entrypoint() -> ComfyExtension: + """ComfyUI calls this at custom-node load time.""" + server_url = (os.environ.get(SERVER_URL_ENV) or "").strip() + if not server_url: + log.info( + "comfy_remote_nodes: %s not set — no remote nodes will be registered", + SERVER_URL_ENV, + ) + return _NullExtension() + log.info("comfy_remote_nodes: connecting to RNP server %s", server_url) + return _RemoteNodesExtension(server_url) + + +class _NullExtension(ComfyExtension): + """Used when no server URL is configured — registers zero nodes.""" + + @override + async def get_node_list(self) -> list[type[IO.ComfyNode]]: + return [] diff --git a/serialization.py b/serialization.py new file mode 100644 index 0000000..990bfa0 --- /dev/null +++ b/serialization.py @@ -0,0 +1,318 @@ +"""Tensor <-> RNP value-envelope helpers. + +The proxy node receives raw ComfyUI Python values from the workflow +runtime (image tensors, mask tensors, audio dicts, Video objects) and +must turn them into wire-format envelopes before POSTing to the RNP +server. Conversely, each ``execute`` response returns envelopes that +must be deserialized into the native ComfyUI types so downstream nodes +in the workflow get exactly what they would have from a local node. + +This module is the single home for that round-trip, so the proxy_node +only ever sees ``encode_*`` / ``decode_*`` pairs and stays focused on +schema construction and HTTP dispatch. + +Byte-level encodings reuse ``comfy_api_nodes.util.conversions`` so a +remote node and a local one produce identical bytes for the same +tensor — no parallel dialect. +""" +from __future__ import annotations + +import base64 +import logging +from io import BytesIO +from typing import Any + +from . import client as rnp_client +from .protocol import ( + Capability, + ErrorCode, + RnpProtocolError, + decode_envelope_data, + is_envelope, +) + +log = logging.getLogger("comfy_remote_nodes.serialization") + + +# --------------------------------------------------------------------------- +# Externalization (inline → presigned-upload swap) +# --------------------------------------------------------------------------- + +# Per-type MIME content-type used when uploading bytes that an envelope +# would otherwise carry inline. The server stamps these on the asset +# record so subsequent fetches arrive with the right Content-Type. +_CONTENT_TYPE_FOR_TYPE = { + "image": "image/png", + "mask": "image/png", + "video": "video/mp4", + "audio": "audio/mpeg", +} + + +async def maybe_externalize( + envelope: dict[str, Any], + *, + server_url: str | None, + max_inline_bytes: int | None, + auth_headers: dict[str, str] | None = None, +) -> dict[str, Any]: + """If the envelope's inline data exceeds ``max_inline_bytes``, + upload the bytes and replace ``data`` with ``uri``. + + No-op when either argument is falsy, the envelope already has a + ``uri``, or the inline payload fits under the cap. + """ + if not server_url or not max_inline_bytes: + return envelope + data = envelope.get("data") + if not isinstance(data, str): + return envelope + # Cap is on raw bytes, not the base64 expansion. + if len(data) * 3 // 4 <= max_inline_bytes: + return envelope + raw = base64.b64decode(data) + content_type = _CONTENT_TYPE_FOR_TYPE.get(envelope.get("type", ""), "application/octet-stream") + download_url = await rnp_client.upload_asset( + server_url, raw, content_type=content_type, auth_headers=auth_headers, + ) + log.info( + "RNP: externalized %s envelope (%d bytes) -> %s", + envelope.get("type"), len(raw), download_url, + ) + out = {k: v for k, v in envelope.items() if k != "data"} + out["uri"] = download_url + return out + + +# --------------------------------------------------------------------------- +# Image encode / decode +# --------------------------------------------------------------------------- + +def encode_image_tensor(tensor: Any) -> dict[str, Any]: + """Encode a ComfyUI IMAGE tensor (B,H,W,C float32 in 0..1) as a + PNG-base64 image envelope. + + Multi-image batches send only the first frame for now; the + framework hasn't grown batch-of-envelopes semantics yet. + """ + import numpy as np + from PIL import Image + + # Strip the batch dim if present. + if hasattr(tensor, "dim") and tensor.dim() == 4: + frame = tensor[0] + else: + frame = tensor + arr = (frame.cpu().clamp(0.0, 1.0).numpy() * 255).astype(np.uint8) + height, width = arr.shape[:2] + channels = arr.shape[2] if arr.ndim == 3 else 1 + pil_mode = {1: "L", 3: "RGB", 4: "RGBA"}.get(channels, "RGB") + img = Image.fromarray(arr, mode=pil_mode) + buf = BytesIO() + img.save(buf, format="PNG") + return { + "type": "image", + "encoding": "png_base64", + "data": base64.b64encode(buf.getvalue()).decode("ascii"), + "shape": [1, height, width, channels], + "dtype": "uint8", + "color_space": "srgb", + } + + +def decode_image_envelope(envelope: dict[str, Any]) -> Any: + """Decode an image envelope into a ComfyUI IMAGE tensor (B,H,W,C float32).""" + encoding = envelope.get("encoding") + if encoding not in ("png_base64", "jpeg_base64", "webp_base64"): + raise RnpProtocolError( + f"Unsupported image encoding: {encoding!r}", + code=ErrorCode.INTERNAL, + ) + import numpy as np + import torch + from PIL import Image + + raw = decode_envelope_data(envelope) + # Server emits PNG (sRGB); RGB is enough for IMAGE — RGBA path + # would force downstream models to receive an alpha channel they + # don't expect. + img = Image.open(BytesIO(raw)).convert("RGB") + arr = np.array(img).astype(np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) + + +# --------------------------------------------------------------------------- +# Mask encode / decode +# --------------------------------------------------------------------------- + +def encode_mask_tensor(tensor: Any) -> dict[str, Any]: + """Encode a ComfyUI MASK tensor (B,H,W float32 in 0..1) as a + PNG-base64 mask envelope (single-channel, 0..255 grayscale).""" + import numpy as np + import torch # noqa: F401 (typing only) + from PIL import Image + + if hasattr(tensor, "dim") and tensor.dim() == 3: + # Batches of masks: take first; multi-mask batches are wrapped + # by the caller as a list of envelopes when needed. + frame = tensor[0] + else: + frame = tensor + arr = (frame.cpu().clamp(0.0, 1.0).numpy() * 255).astype(np.uint8) + img = Image.fromarray(arr, mode="L") + buf = BytesIO() + img.save(buf, format="PNG") + height, width = arr.shape[:2] + return { + "type": "mask", + "encoding": "png_base64", + "data": base64.b64encode(buf.getvalue()).decode("ascii"), + "shape": [1, height, width], + "dtype": "uint8", + } + + +def decode_mask_envelope(envelope: dict[str, Any]) -> Any: + """Decode a mask envelope into a ComfyUI MASK tensor (B,H,W float32).""" + import numpy as np + import torch + from PIL import Image + + encoding = envelope.get("encoding") + if encoding != "png_base64": + raise RnpProtocolError( + f"Unsupported mask encoding: {encoding!r}", + code=ErrorCode.INTERNAL, + ) + raw = decode_envelope_data(envelope) + img = Image.open(BytesIO(raw)).convert("L") + arr = np.array(img).astype(np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) + + +# --------------------------------------------------------------------------- +# Audio encode / decode +# --------------------------------------------------------------------------- + +def encode_audio_input(audio: dict[str, Any]) -> dict[str, Any]: + """Encode a ComfyUI AUDIO dict ({waveform, sample_rate}) as an + mp3-base64 audio envelope. + """ + from comfy_api_nodes.util.conversions import audio_input_to_mp3 + + mp3_buf = audio_input_to_mp3(audio) + sample_rate = int(audio.get("sample_rate", 0)) + waveform = audio.get("waveform") + channels = int(waveform.shape[-2]) if waveform is not None else None + duration_s = ( + float(waveform.shape[-1]) / sample_rate + if (waveform is not None and sample_rate) + else None + ) + extra: dict[str, Any] = {} + if channels is not None: + extra["channels"] = channels + if sample_rate: + extra["sample_rate"] = sample_rate + if duration_s is not None: + extra["duration_s"] = duration_s + return { + "type": "audio", + "encoding": "mp3_base64", + "data": base64.b64encode(mp3_buf.getvalue()).decode("ascii"), + **extra, + } + + +def decode_audio_envelope(envelope: dict[str, Any]) -> dict[str, Any]: + """Decode an audio envelope back into a ComfyUI AUDIO dict.""" + encoding = envelope.get("encoding") + if encoding not in ("mp3_base64", "mp3_inline", "wav_base64", "wav_inline"): + raise RnpProtocolError( + f"Unsupported audio encoding: {encoding!r}", + code=ErrorCode.INTERNAL, + ) + from comfy_api_nodes.util.conversions import audio_bytes_to_audio_input + return audio_bytes_to_audio_input(decode_envelope_data(envelope)) + + +# --------------------------------------------------------------------------- +# Video decode (encode lands when a remote node accepts VIDEO inputs) +# --------------------------------------------------------------------------- + +def decode_video_envelope(envelope: dict[str, Any]) -> Any: + """Decode a video envelope into a ComfyUI Video object (mp4 inline).""" + encoding = envelope.get("encoding") + if encoding not in ("mp4_inline", "mp4_base64"): + raise RnpProtocolError( + f"Unsupported video encoding: {encoding!r}", + code=ErrorCode.INTERNAL, + ) + from comfy_api.latest import InputImpl + return InputImpl.VideoFromFile(BytesIO(decode_envelope_data(envelope))) + + +# --------------------------------------------------------------------------- +# Dispatch by envelope type +# --------------------------------------------------------------------------- + +_DECODERS = { + "image": decode_image_envelope, + "mask": decode_mask_envelope, + "audio": decode_audio_envelope, + "video": decode_video_envelope, +} + + +def decode_envelope(value: dict[str, Any]) -> Any: + """Top-level dispatch: convert any envelope to its native ComfyUI type.""" + decoder = _DECODERS.get(value.get("type", "")) + if decoder is None: + raise RnpProtocolError( + f"Unknown envelope type: {value.get('type')!r}", + code=ErrorCode.INTERNAL, + ) + return decoder(value) + + +def is_image_tensor(value: Any) -> bool: + """Best-effort check for a ComfyUI IMAGE tensor (B,H,W,C float).""" + return _is_torch_tensor(value) and getattr(value, "dim", lambda: 0)() in (3, 4) + + +def is_mask_tensor(value: Any) -> bool: + """Best-effort check for a ComfyUI MASK tensor (B,H,W float or H,W).""" + return _is_torch_tensor(value) and getattr(value, "dim", lambda: 0)() in (2, 3) + + +def is_audio_input(value: Any) -> bool: + return ( + isinstance(value, dict) + and "waveform" in value + and "sample_rate" in value + ) + + +def _is_torch_tensor(value: Any) -> bool: + try: + import torch + except ImportError: + return False + return isinstance(value, torch.Tensor) + + +# Re-export for callers that already imported via this module. +__all__ = [ + "encode_image_tensor", + "decode_image_envelope", + "encode_mask_tensor", + "decode_mask_envelope", + "encode_audio_input", + "decode_audio_envelope", + "decode_video_envelope", + "decode_envelope", + "is_envelope", + "is_image_tensor", + "is_mask_tensor", + "is_audio_input", +]