"""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 asyncio 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 payload exceeds ``max_inline_bytes``, upload it via the presigned-PUT path and rewrite the envelope to reference the uploaded asset(s) instead of inlining base64. Two payload shapes are recognised: * Singleton ``data`` envelopes (the common single-frame IMAGE / MASK / VIDEO / AUDIO case). Oversize payloads upload as one asset and the envelope's ``data`` field is swapped for ``uri``. * Multi-frame ``frames=[…]`` IMAGE envelopes (encoding ``png_base64_batch``). The original ``maybe_externalize`` only inspected ``data`` and so silently let a 9× 2048² Flux2 batch ride the wire as a 200 MB POST body — defeating the entire reason the cap exists. The new path computes the summed raw byte budget (``sum(len(b64)) * 3 // 4`` — base64 expands 4:3 so this is the exact payload the server would receive) and, if it crosses the cap, uploads each frame as an individual presigned PUT in parallel via :func:`asyncio.gather` then rewrites the envelope to ``encoding=png_base64_batch_uri`` with a parallel ``uris`` list. ``shape`` / ``dtype`` / ``color_space`` ride through unchanged so the server-side decoder still sees the batch dimension. Per-frame URIs (NOT a tar/zip blob) compose with the existing per-frame fetch helpers and let the server fetch in parallel. No-op when either argument is falsy, the envelope already has a ``uri`` / ``uris``, or the inline payload fits under the cap. Dispatch keys off the ``encoding`` string — NOT field presence (``frames`` vs. ``data``). An encoder that emitted ``frames`` for a non-batch encoding by accident, or shipped both fields together, used to silently take the batch branch and externalise garbage; keying off the canonical encoding catches that as a hard ``frames`` check inside the batch branch itself. """ if not server_url or not max_inline_bytes: return envelope encoding = envelope.get("encoding") if encoding == "png_base64_batch": frames = envelope.get("frames") if not isinstance(frames, list) or not frames: # Mirror the server-side make_image_batch_envelope contract: # an empty/missing frames list is a hard wire bug, NOT a # silent no-op (which used to ship through under the old # ``isinstance(frames, list) and frames`` dispatch and then # blow up on the server's resolve_image_envelope_pngs check). raise RnpProtocolError( "png_base64_batch envelope requires a non-empty 'frames' list", code=ErrorCode.INTERNAL, ) return await _maybe_externalize_batch( envelope, frames, server_url=server_url, max_inline_bytes=max_inline_bytes, auth_headers=auth_headers, ) 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 async def _maybe_externalize_batch( envelope: dict[str, Any], frames: list[Any], *, server_url: str, max_inline_bytes: int, auth_headers: dict[str, str] | None, ) -> dict[str, Any]: """Externalise a multi-frame IMAGE envelope when its summed inline payload overflows ``max_inline_bytes``. Per-frame upload (NOT a tar/zip blob): each frame becomes its own presigned-PUT asset so the server can fetch in parallel and the existing per-frame ``resolve_image_envelope_pngs`` helper picks up the URI shape transparently. ``asyncio.gather`` issues every PUT in parallel — the dominant latency for a 9-frame Flux2 batch is the upload, not the server hop, so serialising would multiply end-to-end submit time by N. Frame ordering is preserved (gather returns positional results) — critical because a Flux2 provider maps ``uris[k]`` → ``input_image_{k+1}`` 1:1. """ # Cap on raw byte count (base64 expands 4:3, so multiplying by # 3/4 recovers the actual payload the server would receive). total_raw = 0 for f in frames: if isinstance(f, str): total_raw += len(f) * 3 // 4 if total_raw <= max_inline_bytes: return envelope env_type = envelope.get("type", "") content_type = _CONTENT_TYPE_FOR_TYPE.get(env_type, "application/octet-stream") async def _upload_one(idx: int, frame_b64: Any) -> str: if not isinstance(frame_b64, str): raise RnpProtocolError( f"image batch frame {idx} is not a base64 string " f"(got {type(frame_b64).__name__})", code=ErrorCode.INTERNAL, ) raw = base64.b64decode(frame_b64) return await rnp_client.upload_asset( server_url, raw, content_type=content_type, auth_headers=auth_headers, ) uris = await asyncio.gather( *(_upload_one(i, f) for i, f in enumerate(frames)) ) log.info( "RNP: externalized %s batch envelope (%d frames, %d raw bytes) " "-> %d presigned URIs", env_type, len(frames), total_raw, len(uris), ) # Strip the inline ``frames`` field; preserve ``shape`` / ``dtype`` # / ``color_space`` so the server-side decoder still sees the # batch dimension and per-frame metadata. Encoding flips to the # URI variant so dispatch keys cleanly off the string (NOT field # presence — frames vs. uris is a payload field, not a shape # selector). out: dict[str, Any] = { k: v for k, v in envelope.items() if k != "frames" } out["encoding"] = "png_base64_batch_uri" out["uris"] = list(uris) return out # --------------------------------------------------------------------------- # Image encode / decode # --------------------------------------------------------------------------- def encode_image_tensor(tensor: Any, *, accepts_batch: bool = False) -> dict[str, Any]: """Encode a ComfyUI IMAGE tensor (B,H,W,C float32 in 0..1) as an image envelope. ``accepts_batch`` reflects whether the descriptor's ``input_serialization`` for this input lists ``"png_base64_batch"``. When True AND ``tensor.shape[0] > 1``, emit a multi-frame ``png_base64_batch`` envelope (``frames=[, …]``) so the server decodes a real ``[B, H, W, C]`` source. Otherwise emit a single-frame ``png_base64`` envelope for ``tensor[0]`` — preserves backwards compat with descriptors that haven't migrated yet, and matches the wire shape of the existing ``_decode_image_batch_envelope`` on the response path so both directions round-trip symmetrically. """ import numpy as np from PIL import Image rank = getattr(tensor, "dim", lambda: 0)() if rank == 4 and accepts_batch and int(tensor.shape[0]) > 1: frames_arr = (tensor.cpu().clamp(0.0, 1.0).numpy() * 255).astype(np.uint8) b, height, width = frames_arr.shape[:3] channels = frames_arr.shape[3] if frames_arr.ndim == 4 else 1 pil_mode = {1: "L", 3: "RGB", 4: "RGBA"}.get(channels, "RGB") frames_b64: list[str] = [] for i in range(b): img = Image.fromarray(frames_arr[i], mode=pil_mode) buf = BytesIO() img.save(buf, format="PNG") frames_b64.append(base64.b64encode(buf.getvalue()).decode("ascii")) return { "type": "image", "encoding": "png_base64_batch", "frames": frames_b64, "shape": [b, height, width, channels], "dtype": "uint8", "color_space": "srgb", } # Single-frame fallback: matches the old ``encode_image_tensor`` # contract (sends ``tensor[0]`` for a multi-frame batch when the # descriptor doesn't advertise the batch encoding). if rank == 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. Single-frame encodings (``png_base64``, ``jpeg_base64``, ``webp_base64``) decode into ``[1, H, W, C]``. The multi-frame ``png_base64_batch`` encoding (a list of base64-PNG ``frames``) stacks each frame into ``[B, H, W, C]`` — the native ComfyUI image-batch shape — so providers whose upstream returns ``n > 1`` surface as a real batched tensor in the workflow runtime. """ encoding = envelope.get("encoding") if encoding == "png_base64_batch": return _decode_image_batch_envelope(envelope) 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) def _decode_image_batch_envelope(envelope: dict[str, Any]) -> Any: """Decode a ``png_base64_batch`` envelope into ``[B, H, W, C]``. The wire shape is ``{"type":"image", "encoding":"png_base64_batch", "frames":[, ...], "shape":[B,H,W,C], ...}``. Each frame is a PNG-encoded image; all frames must share the same H/W after the RGB conversion (the server-side encoder enforces this). A frame count mismatch against the ``shape[0]`` field surfaces as a hard parse error rather than silently dropping frames — matches the SVG decoder's ``count`` integrity check. """ import base64 as _b64 import numpy as np import torch from PIL import Image frames = envelope.get("frames") if not isinstance(frames, list) or not frames: raise RnpProtocolError( "image batch envelope missing or empty 'frames' list", code=ErrorCode.INTERNAL, ) declared_shape = envelope.get("shape") if isinstance(declared_shape, list) and len(declared_shape) >= 1: declared_b = int(declared_shape[0]) if declared_b != len(frames): raise RnpProtocolError( f"image batch envelope frame count {len(frames)} does not " f"match declared shape[0]={declared_b}", code=ErrorCode.INTERNAL, ) arrs: list[Any] = [] base_hw: tuple[int, int] | None = None for idx, frame_b64 in enumerate(frames): if not isinstance(frame_b64, str): raise RnpProtocolError( f"image batch frame {idx} is not a base64 string", code=ErrorCode.INTERNAL, ) try: raw = _b64.b64decode(frame_b64) img = Image.open(BytesIO(raw)).convert("RGB") except Exception as e: raise RnpProtocolError( f"image batch frame {idx} failed to decode: {e}", code=ErrorCode.INTERNAL, ) from e arr = np.array(img).astype(np.float32) / 255.0 if base_hw is None: base_hw = (arr.shape[0], arr.shape[1]) elif (arr.shape[0], arr.shape[1]) != base_hw: raise RnpProtocolError( f"image batch frame {idx} size {arr.shape[:2]} does not " f"match frame 0 size {base_hw} (server should have rejected " f"this; likely a wire bug)", code=ErrorCode.INTERNAL, ) arrs.append(arr) stacked = np.stack(arrs, axis=0) return torch.from_numpy(stacked) # --------------------------------------------------------------------------- # 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", "opus_base64", ): 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))) # --------------------------------------------------------------------------- # 3D model decode (encode lands when a remote node accepts MODEL_3D inputs) # --------------------------------------------------------------------------- # --------------------------------------------------------------------------- # Multi-file 3D bundle: a ``File3D``-compatible class that carries the # primary mesh plus N companion files (textures / .mtl / .bin / etc.). # # Built lazily as a ``File3D`` subclass on first use — importing # ``File3D`` at module load pulls torch (via ``comfy_api.latest._util. # geometry_types``), which is fine inside ComfyUI but breaks the # smoke-test pattern that exercises serialization without the heavy # comfy_api tree. The factory below caches the constructed class so # every ``BundledFile3D`` instance shares the same dynamic class and # ``isinstance(x, File3D)`` always succeeds. # --------------------------------------------------------------------------- _BUNDLED_FILE3D_CLS: Any = None def _bundled_file3d_class() -> type: """Return the cached ``BundledFile3D`` class (a ``File3D`` subclass). Constructs it on first call so the ``comfy_api.latest._util`` import (which pulls torch) is paid for only when a bundle envelope is actually decoded — keeps the smoke-test stubbing pattern intact for callers that don't exercise this path. """ global _BUNDLED_FILE3D_CLS if _BUNDLED_FILE3D_CLS is not None: return _BUNDLED_FILE3D_CLS from comfy_api.latest._util.geometry_types import File3D as _File3D class BundledFile3D(_File3D): """``File3D``-compatible wrapper for a multi-file 3D bundle. Stores the primary mesh file plus N companion files (textures / .mtl / .bin / etc.) in memory as a ``path -> bytes`` dict; lazily materialises the whole bundle into a temp directory the first time anything cross-file-resolution-sensitive is asked of it (``get_source()`` for an OBJ/GLTF that needs its .mtl/.bin/textures on disk; ``save_to(dest)`` which copies the *whole bundle* into the destination directory). Cheap consumers that just want the primary file's bytes (``get_data()`` / ``get_bytes()``) skip the temp-dir materialise entirely. The temp directory's lifetime is tied to this object: it is cleaned up on garbage-collect via ``__del__`` (best-effort). For workflows that need stable on-disk paths beyond this object's lifetime, call ``save_to(dest)`` to copy the bundle into a caller-controlled location. """ def __init__( self, files: dict[str, bytes], primary_path: str, roles: dict[str, str] | None = None, formats: dict[str, str] | None = None, ) -> None: if primary_path not in files: raise ValueError( f"primary_path {primary_path!r} not in bundle files: " f"{sorted(files)!r}" ) self._files: dict[str, bytes] = dict(files) self._primary_path = primary_path self._roles: dict[str, str] = dict(roles or {}) self._formats: dict[str, str] = dict(formats or {}) self._temp_dir: str | None = None # Populate parent ``_source`` / ``_format`` so reads of # those plain attributes on downstream consumers that # haven't been updated still work. ``_source`` is a # BytesIO of the primary file; ``_format`` is the primary # file's extension. primary_fmt = ( self._formats.get(primary_path) or ( primary_path.rsplit(".", 1)[-1].lower() if "." in primary_path else "" ) ) super().__init__(BytesIO(self._files[primary_path]), primary_fmt) # -------------------------------------------------------------- # File3D-compatible accessors (overrides) # -------------------------------------------------------------- def get_source(self) -> str: """Return the on-disk path to the primary file. Materialises the whole bundle into a temp directory on first call so cross-file refs (.obj -> .mtl -> texture; .gltf -> .bin + textures) resolve via the filesystem. """ self._ensure_materialized() from os.path import join as _join return _join(self._temp_dir, *self._primary_path.split("/")) def get_data(self) -> BytesIO: """Return the primary file's bytes (no disk I/O).""" return BytesIO(self._files[self._primary_path]) def get_bytes(self) -> bytes: """Return the primary file's raw bytes (no disk I/O).""" return self._files[self._primary_path] def save_to(self, path: str) -> str: """Copy the WHOLE bundle into the destination directory. ``path`` is the destination for the *primary file*; sidecars are copied alongside it preserving their relative paths from the original bundle layout. E.g. for primary_path = "model.obj" files = { "model.obj": , "model.mtl": , "textures/albedo.png": , } ``save_to("/tmp/out/scene.obj")`` writes: /tmp/out/scene.obj /tmp/out/model.mtl /tmp/out/textures/albedo.png Sidecars keep their ORIGINAL filenames (so cross-file refs inside the .obj that point at ``model.mtl`` still resolve); only the primary file is renamed to match ``path``. The destination directory is created if needed. """ from pathlib import Path dest = Path(path) dest.parent.mkdir(parents=True, exist_ok=True) dest.write_bytes(self._files[self._primary_path]) dest_parent_resolved = dest.parent.resolve() for rel_path, raw in self._files.items(): if rel_path == self._primary_path: continue sidecar_dest = dest.parent.joinpath(*rel_path.split("/")) # Path safety — rel_path was sanitized at envelope # build time, but defence-in-depth check on the # resolved path keeps us safe against pathological # inputs that may have slipped through. resolved = sidecar_dest.resolve() if dest_parent_resolved not in resolved.parents and resolved != dest_parent_resolved: continue sidecar_dest.parent.mkdir(parents=True, exist_ok=True) sidecar_dest.write_bytes(raw) return str(dest) # -------------------------------------------------------------- # Bundle-specific accessors (additive) # -------------------------------------------------------------- @property def primary_path(self) -> str: return self._primary_path @property def file_paths(self) -> list[str]: return sorted(self._files) def file_bytes(self, path: str) -> bytes: """Raw bytes for one file in the bundle by relative path.""" return self._files[path] def file_role(self, path: str) -> str | None: """Optional producer-supplied role hint (may be ``None``).""" return self._roles.get(path) def paths_with_role(self, role: str) -> list[str]: """All file paths whose role hint equals ``role``.""" return [p for p, r in self._roles.items() if r == role] def materialize_to_temp_dir(self) -> str: """Force-materialise the whole bundle into a temp directory. Returns the temp-directory path. Cleanup happens on GC; callers that need stable lifetimes should ``save_to()`` to a caller-controlled location instead. """ self._ensure_materialized() return self._temp_dir # -------------------------------------------------------------- # Internal # -------------------------------------------------------------- def _ensure_materialized(self) -> None: if self._temp_dir is not None: return import tempfile from pathlib import Path self._temp_dir = tempfile.mkdtemp(prefix="rnp_model3d_bundle_") root = Path(self._temp_dir) for rel_path, raw in self._files.items(): # rel_path was sanitized at envelope build time — # split-on-'/' and re-join via Path is safe. dest = root.joinpath(*rel_path.split("/")) dest.parent.mkdir(parents=True, exist_ok=True) dest.write_bytes(raw) def __del__(self) -> None: td = getattr(self, "_temp_dir", None) if td is not None: try: import shutil shutil.rmtree(td, ignore_errors=True) except Exception: # noqa: BLE001 pass def __repr__(self) -> str: return ( f"BundledFile3D(primary={self._primary_path!r}, " f"files={len(self._files)}, format={self._format!r})" ) _BUNDLED_FILE3D_CLS = BundledFile3D return _BUNDLED_FILE3D_CLS def decode_model3d_envelope(envelope: dict[str, Any]) -> Any: """Decode a 3D-model envelope into a ComfyUI ``File3D`` object. Dispatches on the envelope's ``encoding``: * ``glb_inline`` — single-file binary glTF 2.0 (magic ``b"glTF"`` at offset 0). The bytes are handed to ``File3D`` as a ``BytesIO`` so downstream nodes (preview3d / SaveGLB / Load3D) see exactly the same value shape they would from a local Tripo / Rodin / Meshy node. The envelope's optional ``format`` extra is forwarded as the second constructor arg so the materialiser picks the right file extension when persisting. * ``bundle_inline`` — multi-file bundle. The primary mesh file plus N companion files (textures / .mtl / .bin / etc.) are decoded into a :class:`BundledFile3D` (``File3D``-compatible) whose ``get_source()`` / ``save_to()`` materialise the whole bundle on disk so cross-file refs resolve via the filesystem. Path safety (no ``..`` / no absolute prefix / no Windows drive letter / no case-insensitive collisions) was enforced at envelope build time on the server; the client re-asserts the ``primary_path``-present invariant defensively and rejects any file entry that lacks ``data``. """ encoding = envelope.get("encoding") if encoding == "glb_inline": from comfy_api.latest._util import File3D file_format = envelope.get("format") or "glb" return File3D(BytesIO(decode_envelope_data(envelope)), str(file_format)) if encoding == "bundle_inline": bundle_cls = _bundled_file3d_class() primary_path = envelope.get("primary_path") if not isinstance(primary_path, str) or not primary_path: raise RnpProtocolError( "bundle_inline envelope missing 'primary_path'", code=ErrorCode.INTERNAL, ) files_spec = envelope.get("files") if not isinstance(files_spec, list) or not files_spec: raise RnpProtocolError( "bundle_inline envelope missing non-empty 'files' list", code=ErrorCode.INTERNAL, ) files: dict[str, bytes] = {} roles: dict[str, str] = {} formats: dict[str, str] = {} for entry in files_spec: if not isinstance(entry, dict): raise RnpProtocolError( f"bundle_inline file entry not a dict: {entry!r}", code=ErrorCode.INTERNAL, ) path = entry.get("path") if not isinstance(path, str) or not path: raise RnpProtocolError( f"bundle_inline file entry missing 'path': {entry!r}", code=ErrorCode.INTERNAL, ) data = entry.get("data") if not isinstance(data, str): raise RnpProtocolError( f"bundle_inline file {path!r} missing inline 'data'", code=ErrorCode.INTERNAL, ) files[path] = base64.b64decode(data) role = entry.get("role") if isinstance(role, str): roles[path] = role fmt = entry.get("format") if isinstance(fmt, str): formats[path] = fmt if primary_path not in files: raise RnpProtocolError( f"bundle_inline primary_path {primary_path!r} not in " f"files list (have: {sorted(files)!r})", code=ErrorCode.INTERNAL, ) return bundle_cls( files=files, primary_path=primary_path, roles=roles, formats=formats, ) raise RnpProtocolError( f"Unsupported model_3d encoding: {encoding!r}", code=ErrorCode.INTERNAL, ) # --------------------------------------------------------------------------- # Cross-node task handle # # A ``task_handle`` envelope wraps a vendor-native task / video / job # ID (Tripo MODEL_TASK_ID, Kling video_id, ...) with lineage metadata # so chained provider nodes can pass the handle through the workflow # without losing the upstream chain. The decoder returns a # :class:`TaskHandle` dataclass — re-encoding the same dataclass on a # downstream submit (via :func:`encode_task_handle`) preserves the # full ``parent_chain`` without the proxy_node having to know # anything provider-specific. Plain ``str`` round-trip would silently # drop lineage, so the decoder deliberately surfaces the rich type # even at the cost of slightly more work on the consumer side. # --------------------------------------------------------------------------- from dataclasses import dataclass, field @dataclass class TaskHandle: """Cross-node handle to a vendor-native async task / video / job. See :data:`comfy_remote_nodes.protocol.Capability.IO_TASK_HANDLE` for the wire-shape contract. The dataclass mirrors the envelope fields verbatim so :func:`encode_task_handle` is a trivial round-trip. ``parent_chain`` is a flat list of ``{vendor, kind, native_id, origin_node_id}`` dicts in root-to-immediate-parent order; treat it as opaque on the client — the server is the only component that ever consumes it. """ vendor: str kind: str native_id: str origin_node_id: str parent_chain: list[dict[str, str]] = field(default_factory=list) def __str__(self) -> str: # Workflow runtime sometimes coerces socket values to ``str`` # (UI previews, log lines). Surfacing the native id keeps that # output human-meaningful without leaking the lineage chain # into UI strings. return self.native_id def _validate_task_handle_ref(ref: Any, *, context: str) -> dict[str, str]: """Validate one ``{vendor, kind, native_id, origin_node_id}`` ref. Mirror of the server-side ``_validate_task_handle_ref`` — rejects malformed parent_chain entries with the same field-by- field error messages so a bad envelope surfaces the same way on both sides. """ if not isinstance(ref, dict): raise RnpProtocolError( f"{context}: expected dict, got {type(ref).__name__}", code=ErrorCode.INTERNAL, ) out: dict[str, str] = {} for fld in ("vendor", "kind", "native_id", "origin_node_id"): val = ref.get(fld) if not isinstance(val, str) or not val: raise RnpProtocolError( f"{context}: field {fld!r} must be a non-empty string " f"(got {val!r})", code=ErrorCode.INTERNAL, ) out[fld] = val if "," in out["kind"]: raise RnpProtocolError( f"{context}: 'kind' must be a singular socket kind, not a " f"comma-union string (got {out['kind']!r})", code=ErrorCode.INTERNAL, ) return out def decode_task_handle_envelope(envelope: dict[str, Any]) -> TaskHandle: """Decode a ``task_handle`` envelope into a :class:`TaskHandle`. Only the ``vendor_inline`` encoding is recognised today (a future sibling ``replay_ticket_inline`` would carry a signed server- local replay token for cross-server portability — not in this PR). All four top-level fields and every ``parent_chain`` entry are validated; rejections raise ``RnpProtocolError(INTERNAL)`` with the same field-by-field messages the server emits at envelope-build time. """ encoding = envelope.get("encoding") if encoding != "vendor_inline": raise RnpProtocolError( f"Unsupported task_handle encoding: {encoding!r}", code=ErrorCode.INTERNAL, ) top = _validate_task_handle_ref( { "vendor": envelope.get("vendor"), "kind": envelope.get("kind"), "native_id": envelope.get("native_id"), "origin_node_id": envelope.get("origin_node_id"), }, context="task_handle", ) chain_in = envelope.get("parent_chain", []) if not isinstance(chain_in, list): raise RnpProtocolError( f"task_handle 'parent_chain' must be a list (got " f"{type(chain_in).__name__})", code=ErrorCode.INTERNAL, ) chain_out = [ _validate_task_handle_ref(entry, context=f"parent_chain[{i}]") for i, entry in enumerate(chain_in) ] return TaskHandle( vendor=top["vendor"], kind=top["kind"], native_id=top["native_id"], origin_node_id=top["origin_node_id"], parent_chain=chain_out, ) def encode_task_handle(handle: TaskHandle) -> dict[str, Any]: """Encode a :class:`TaskHandle` back into a wire-format envelope. Symmetric counterpart of :func:`decode_task_handle_envelope` so a handle decoded from one provider's output and fed into another provider's input round-trips losslessly — the proxy_node calls this whenever it sees a ``TaskHandle`` in the input values for a descriptor that accepts the ``task_handle`` envelope type. The returned dict is the same shape the server's ``make_task_handle_envelope`` produces, so the server-side decoder accepts it without any branching on "is this a fresh handle or a re-encoded one" — the wire is the source of truth. """ return { "type": "task_handle", "encoding": "vendor_inline", "vendor": handle.vendor, "kind": handle.kind, "native_id": handle.native_id, "origin_node_id": handle.origin_node_id, "parent_chain": [dict(entry) for entry in handle.parent_chain], } # --------------------------------------------------------------------------- # Dispatch by envelope type # --------------------------------------------------------------------------- _DECODERS = { "image": decode_image_envelope, "mask": decode_mask_envelope, "audio": decode_audio_envelope, "video": decode_video_envelope, "model_3d": decode_model3d_envelope, "task_handle": decode_task_handle_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_model3d_envelope", "_bundled_file3d_class", "TaskHandle", "decode_task_handle_envelope", "encode_task_handle", "decode_envelope", "is_envelope", "is_image_tensor", "is_mask_tensor", "is_audio_input", ]