From ed63c5ae49d639eb43458655b9c22185290eb036 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 13 May 2026 22:00:52 -0700 Subject: [PATCH] feat(serialization): decode image:png_base64_batch envelope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the server-side change in comfy-rnp-server PR #7. Adds the client-side half of the multi-image envelope contract: - Capability.IMAGE_PNG_BASE64_BATCH (\image:png_base64_batch\) added to protocol.Capability, mirroring comfy_rnp_protocol/constants.py. - decode_image_envelope() now dispatches png_base64_batch to a new _decode_image_batch_envelope() helper that stacks each PNG frame into a [B, H, W, C] tensor — the native ComfyUI image-batch shape — so providers whose upstream returns n>1 surface as a real batched tensor in the workflow runtime instead of being silently truncated to the first frame. - Frame-count integrity check: if the envelope declares shape[0]=B but \ rames\ has a different length, fail hard rather than drop frames silently. Matches the SVG decoder's \count\ check. - Per-frame H/W mismatch is also a hard error — the server-side encoder enforces uniform sizing already, so a mismatch on the wire indicates a wire bug worth surfacing. - IMAGE_PNG_BASE64_BATCH added to client.CLIENT_CAPABILITIES so the server's per-request negotiator (\client_advertises_capability\ in comfy_rnp_protocol.capabilities) lights up the batch path. Single-frame envelopes (\png_base64\ / \jpeg_base64\ / \webp_base64\) remain unchanged — back-compat with every already-merged provider. Verified end-to-end with a synthetic 2-frame envelope (white + black 32x64 PNGs) plus the mismatch / empty-frames error paths. Co-authored-by: Amp Amp-Thread-ID: https://ampcode.com/threads/T-019e24d1-250c-750a-ab00-2eb5e9cb2500 --- client.py | 10 +++++-- protocol.py | 6 ++++ serialization.py | 75 +++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 88 insertions(+), 3 deletions(-) diff --git a/client.py b/client.py index a7895e9..fb3a5a6 100644 --- a/client.py +++ b/client.py @@ -30,12 +30,18 @@ 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. +# Gemini, Ideogram, ...) need, plus the ``schema:v3`` marker so +# descriptors come back in the V3 ``get_v1_info()`` shape. Servers use +# this list to negotiate per-request: e.g. a provider only emits the +# multi-frame ``image:png_base64_batch`` envelope when this list +# contains :data:`Capability.IMAGE_PNG_BASE64_BATCH`, otherwise it +# falls back to the single-frame ``image:png_base64`` shape (with a +# server-side truncation warning logged when extras are dropped). CLIENT_CAPABILITIES = [ Capability.SCHEMA_V3, Capability.VIDEO_MP4_INLINE, Capability.IMAGE_PNG_BASE64, + Capability.IMAGE_PNG_BASE64_BATCH, Capability.MASK_PNG_BASE64, ] diff --git a/protocol.py b/protocol.py index 73971e0..66d263d 100644 --- a/protocol.py +++ b/protocol.py @@ -93,6 +93,12 @@ class Capability: IMAGE_PNG_BASE64 = "image:png_base64" IMAGE_JPEG_BASE64 = "image:jpeg_base64" IMAGE_RAW_CHW_UINT8 = "image:raw_b64_chw_uint8" + # Multi-frame IMAGE envelope: a single value carries N PNG-base64 + # frames sharing the same H/W/C, decoded into a [B,H,W,C] tensor on + # this client. Mirrors the AUDIO sweep precedent (parallel encodings + # rather than extending ``png_base64``) so legacy server emits still + # decode correctly when the new token isn't advertised. + IMAGE_PNG_BASE64_BATCH = "image:png_base64_batch" VIDEO_MP4_INLINE = "video:mp4_inline" AUDIO_MP3_INLINE = "audio:mp3_inline" MASK_PNG_BASE64 = "mask:png_base64" diff --git a/serialization.py b/serialization.py index 990bfa0..a77a93f 100644 --- a/serialization.py +++ b/serialization.py @@ -121,8 +121,18 @@ def encode_image_tensor(tensor: Any) -> dict[str, Any]: def decode_image_envelope(envelope: dict[str, Any]) -> Any: - """Decode an image envelope into a ComfyUI IMAGE tensor (B,H,W,C float32).""" + """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}", @@ -141,6 +151,69 @@ def decode_image_envelope(envelope: dict[str, Any]) -> Any: 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 # ---------------------------------------------------------------------------