Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0946b83b21 | ||
|
|
866b7e32a3 | ||
|
|
1d70ef223b | ||
|
|
197e7ba406 | ||
|
|
95a3073d79 | ||
|
|
01f71bd806 | ||
|
|
569212f1e7 | ||
|
|
04ca70a853 | ||
|
|
ed5823e8e7 | ||
|
|
22e0742d47 | ||
|
|
90ecd3218c |
@@ -103,20 +103,6 @@ CLIENT_CAPABILITIES = [
|
||||
# this token; legacy clients surface NEGOTIATION_FAILED at
|
||||
# descriptor-load time rather than crashing at execute.
|
||||
Capability.MODEL_3D_BUNDLE_INLINE,
|
||||
# Cross-node task handle envelope (vendor-native task / video /
|
||||
# job IDs with lineage metadata for chained partner provider
|
||||
# nodes — Tripo Texture/Refine/Rig/Retarget/Conversion, Kling
|
||||
# VideoExtend, future partner chains). The decoder in
|
||||
# ``serialization.decode_task_handle_envelope`` returns a
|
||||
# ``TaskHandle`` dataclass rather than a plain ``str`` so the
|
||||
# ``parent_chain`` lineage round-trips when the handle is passed
|
||||
# through to a downstream provider node (re-encoded by
|
||||
# ``serialization.encode_task_handle`` on the submit path).
|
||||
# Servers gate handle-emitting AND handle-accepting descriptors
|
||||
# on this token (whole-descriptor gating — partial-socket
|
||||
# synthesis is rejected because a chain that breaks meaningfully
|
||||
# if either end is missing is worse than no chain at all).
|
||||
Capability.IO_TASK_HANDLE,
|
||||
]
|
||||
|
||||
# Reported in the X-RNP-Client-Version header.
|
||||
|
||||
@@ -0,0 +1,360 @@
|
||||
"""Smoke test for nested dict recursion in ``_encode_one`` /
|
||||
``_encode_inputs`` / ``_externalize_nested``.
|
||||
|
||||
Background: ComfyUI's V3 ``build_nested_inputs`` rebuilds AUTOGROW and
|
||||
DYNAMIC_COMBO runtime kwargs into nested dicts (e.g. ``{"reference_images":
|
||||
{"image1": <tensor>, "image2": <tensor>}}`` for a top-level AG, or
|
||||
``{"model": {"branch": "wan2.7-r2v", "video_refs": {"video1":
|
||||
<VideoInput>}}}`` for AG-inside-DC). Before this fix, ``_encode_one``
|
||||
returned a non-tensor non-VideoInput non-audio dict unchanged, which
|
||||
meant per-slot tensors / VideoInputs survived past encoding and would
|
||||
either crash at ``json.dumps`` time or violate the server's contract
|
||||
(comfy_rnp_server's ``_ordered_wan_videoedit_image_envelopes`` etc.
|
||||
expect each slot value to already be an RNP envelope).
|
||||
|
||||
Asserts:
|
||||
|
||||
1. Top-level AG-of-IMAGE dict: each ``image<N>`` tensor becomes an
|
||||
image envelope; the wrapping dict structure is preserved.
|
||||
2. Top-level AG-of-VIDEO dict: each ``video<N>`` VideoInput becomes a
|
||||
video envelope.
|
||||
3. DC-wrapping-AG: top-level dict carries a branch key (string) AND a
|
||||
nested AG dict of tensors / VideoInputs; encoder recurses one level
|
||||
deeper and encodes the leaves while leaving the branch key untouched.
|
||||
4. AUDIO input is NOT swallowed by the dict-recursion branch: the
|
||||
``{"waveform": ..., "sample_rate": ...}`` shape still routes to
|
||||
``encode_audio_input``.
|
||||
5. Already-encoded envelope dicts pass through unchanged (no double-
|
||||
encoding) — protects against re-running the encoder on a value that
|
||||
the caller pre-encoded.
|
||||
6. Plain scalar config dicts (``{"width": 1024, "height": 576}``) pass
|
||||
through unchanged.
|
||||
7. ``_externalize_nested`` walks the same shape: a nested AG dict of
|
||||
oversize envelopes uploads each envelope via the stubbed externalize
|
||||
path and the envelope's ``data`` field is swapped for ``uri``;
|
||||
non-envelope siblings are untouched; the wrapping dict structure is
|
||||
preserved.
|
||||
8. JSON-serializability post-condition: ``json.dumps`` round-trips the
|
||||
full encoded payload (regression guard for the exact bug fixed).
|
||||
9. Single-tensor / single-VideoInput / single-audio top-level inputs
|
||||
still encode the same way they did before (no regression for the
|
||||
existing leaf path).
|
||||
|
||||
Run with any python that has ``torch`` available (uses the ComfyUI
|
||||
venv on this workstation):
|
||||
|
||||
python notes/run_encode_nested_dicts.py
|
||||
|
||||
The test loads ``_encode_one`` / ``_encode_inputs`` / ``_externalize_nested``
|
||||
out of ``proxy_node.py`` via AST extraction so we don't pay the cost
|
||||
of importing ComfyUI's ``comfy_api`` / ``comfy_api_nodes`` packages
|
||||
(which those helpers don't touch). Same standalone pattern as
|
||||
``run_image_max_batch.py``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
import torch
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Locate the worktree under test (this file's parent's parent).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
_CLIENT_DIR = os.path.dirname(_HERE)
|
||||
|
||||
|
||||
def _load_protocol() -> types.ModuleType:
|
||||
"""Load ``protocol.py`` standalone (no relative imports)."""
|
||||
path = os.path.join(_CLIENT_DIR, "protocol.py")
|
||||
spec = importlib.util.spec_from_file_location("rnp_smoke_protocol", path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
class _FakeVideoInput:
|
||||
"""Duck-typed VideoInput stand-in: ``save_to`` + ``get_duration``
|
||||
are the only attrs ``serialization.is_video_input`` checks for."""
|
||||
def __init__(self, tag: str, duration: float = 4.0) -> None:
|
||||
self.tag = tag
|
||||
self._duration = duration
|
||||
|
||||
def save_to(self, *args, **kwargs): # pragma: no cover — never called by stub
|
||||
return None
|
||||
|
||||
def get_duration(self) -> float:
|
||||
return self._duration
|
||||
|
||||
|
||||
def _build_serialization_stub(protocol: types.ModuleType) -> types.ModuleType:
|
||||
"""Minimal ``serialization`` shim covering exactly the surface the
|
||||
helpers under test use."""
|
||||
mod = types.ModuleType("rnp_smoke_serialization")
|
||||
|
||||
def _is_torch_tensor(value): # noqa: ANN001
|
||||
return isinstance(value, torch.Tensor)
|
||||
|
||||
def is_audio_input(value): # noqa: ANN001
|
||||
return (
|
||||
isinstance(value, dict)
|
||||
and {"waveform", "sample_rate"} <= set(value.keys())
|
||||
)
|
||||
|
||||
def is_video_input(value): # noqa: ANN001
|
||||
return (
|
||||
callable(getattr(value, "save_to", None))
|
||||
and callable(getattr(value, "get_duration", None))
|
||||
)
|
||||
|
||||
def is_model3d_input(value): # noqa: ANN001
|
||||
return False # unused in this test
|
||||
|
||||
def encode_image_tensor(tensor, *, accepts_batch=False): # noqa: ANN001, ARG001
|
||||
return {
|
||||
"type": "image",
|
||||
"encoding": "png_base64",
|
||||
"data": f"img:{int(tensor.shape[0])}",
|
||||
}
|
||||
|
||||
def encode_mask_tensor(tensor): # noqa: ANN001, ARG001
|
||||
return {"type": "mask", "encoding": "png_base64", "data": "mask"}
|
||||
|
||||
def encode_audio_input(value): # noqa: ANN001, ARG001
|
||||
return {"type": "audio", "encoding": "mp3_base64", "data": "aud"}
|
||||
|
||||
def encode_video_input(value): # noqa: ANN001
|
||||
return {
|
||||
"type": "video",
|
||||
"encoding": "mp4_base64",
|
||||
"data": f"vid:{getattr(value, 'tag', '?')}",
|
||||
}
|
||||
|
||||
def encode_model3d_input(value): # noqa: ANN001, ARG001 # pragma: no cover
|
||||
return {"type": "model_3d", "encoding": "glb_base64", "data": "m3d"}
|
||||
|
||||
HEAVY_TYPES = {"image", "mask", "audio", "video", "model_3d"}
|
||||
|
||||
def is_envelope(value): # noqa: ANN001
|
||||
return (
|
||||
isinstance(value, dict)
|
||||
and isinstance(value.get("type"), str)
|
||||
and value.get("type") in HEAVY_TYPES
|
||||
and isinstance(value.get("encoding"), str)
|
||||
)
|
||||
|
||||
# Cap-aware ``maybe_externalize`` stub: swap ``data`` for ``uri`` when
|
||||
# the inline payload would exceed ``max_inline_bytes``. Lets us prove
|
||||
# that ``_externalize_nested`` reaches every per-slot envelope.
|
||||
upload_log: list[str] = []
|
||||
|
||||
async def maybe_externalize(
|
||||
envelope, # noqa: ANN001
|
||||
*, server_url=None, max_inline_bytes=None, auth_headers=None, # noqa: ANN001, ARG001
|
||||
):
|
||||
if not server_url or max_inline_bytes is None:
|
||||
return envelope
|
||||
data = envelope.get("data")
|
||||
if not isinstance(data, str):
|
||||
return envelope
|
||||
if len(data) <= max_inline_bytes:
|
||||
return envelope
|
||||
upload_log.append(f"{envelope.get('type')}:{data}")
|
||||
out = {k: v for k, v in envelope.items() if k != "data"}
|
||||
out["uri"] = f"https://upload/{envelope.get('type')}/{len(upload_log)}"
|
||||
return out
|
||||
|
||||
mod._is_torch_tensor = _is_torch_tensor
|
||||
mod.is_audio_input = is_audio_input
|
||||
mod.is_video_input = is_video_input
|
||||
mod.is_model3d_input = is_model3d_input
|
||||
mod.encode_image_tensor = encode_image_tensor
|
||||
mod.encode_mask_tensor = encode_mask_tensor
|
||||
mod.encode_audio_input = encode_audio_input
|
||||
mod.encode_video_input = encode_video_input
|
||||
mod.encode_model3d_input = encode_model3d_input
|
||||
mod.is_envelope = is_envelope
|
||||
mod.maybe_externalize = maybe_externalize
|
||||
mod._upload_log = upload_log # exposed for assertions
|
||||
return mod
|
||||
|
||||
|
||||
def _build_proxy_node_stub(
|
||||
protocol: types.ModuleType, serialization: types.ModuleType,
|
||||
) -> types.ModuleType:
|
||||
"""Extract the helpers under test from ``proxy_node.py``."""
|
||||
mod = types.ModuleType("rnp_smoke_proxy_node")
|
||||
mod.__dict__.update({
|
||||
"Any": object,
|
||||
"log": types.SimpleNamespace(warning=lambda *a, **kw: None),
|
||||
"serialization": serialization,
|
||||
"RnpProtocolError": protocol.RnpProtocolError,
|
||||
"ErrorCode": protocol.ErrorCode,
|
||||
})
|
||||
src_path = os.path.join(_CLIENT_DIR, "proxy_node.py")
|
||||
with open(src_path, "r", encoding="utf-8") as fh:
|
||||
src = fh.read()
|
||||
tree = ast.parse(src)
|
||||
wanted = {
|
||||
"_enforce_local_validate",
|
||||
"_check_image_max_batch",
|
||||
"_encode_one",
|
||||
"_encode_inputs",
|
||||
"_externalize_nested",
|
||||
}
|
||||
nodes = [
|
||||
n for n in tree.body
|
||||
if isinstance(n, (ast.AsyncFunctionDef, ast.FunctionDef))
|
||||
and n.name in wanted
|
||||
]
|
||||
snippet = ast.Module(body=nodes, type_ignores=[])
|
||||
code = compile(snippet, src_path, "exec")
|
||||
exec(code, mod.__dict__)
|
||||
missing = wanted - set(mod.__dict__)
|
||||
if missing:
|
||||
raise RuntimeError(f"proxy_node helpers missing from extract: {missing}")
|
||||
return mod
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def main() -> int:
|
||||
protocol = _load_protocol()
|
||||
serialization = _build_serialization_stub(protocol)
|
||||
proxy_node = _build_proxy_node_stub(protocol, serialization)
|
||||
|
||||
img1 = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
img2 = torch.zeros((1, 32, 32, 3), dtype=torch.float32)
|
||||
vid1 = _FakeVideoInput("v1")
|
||||
vid2 = _FakeVideoInput("v2")
|
||||
audio = {
|
||||
"waveform": torch.zeros((1, 1, 1000), dtype=torch.float32),
|
||||
"sample_rate": 16000,
|
||||
}
|
||||
|
||||
# ---- 1. Top-level AG-of-IMAGE dict encodes per-slot
|
||||
out = await proxy_node._encode_inputs(
|
||||
{"reference_images": {"image1": img1, "image2": img2}},
|
||||
)
|
||||
refs = out["reference_images"]
|
||||
assert isinstance(refs, dict) and set(refs.keys()) == {"image1", "image2"}, refs
|
||||
assert refs["image1"]["type"] == "image" and refs["image1"]["data"] == "img:1"
|
||||
assert refs["image2"]["type"] == "image" and refs["image2"]["data"] == "img:1"
|
||||
print("ok: top-level AG-of-IMAGE encoded per slot")
|
||||
|
||||
# ---- 2. Top-level AG-of-VIDEO dict encodes per-slot
|
||||
out = await proxy_node._encode_inputs(
|
||||
{"video_refs": {"video1": vid1, "video2": vid2}},
|
||||
)
|
||||
vrefs = out["video_refs"]
|
||||
assert isinstance(vrefs, dict) and set(vrefs.keys()) == {"video1", "video2"}, vrefs
|
||||
assert vrefs["video1"]["type"] == "video" and vrefs["video1"]["data"] == "vid:v1"
|
||||
assert vrefs["video2"]["type"] == "video" and vrefs["video2"]["data"] == "vid:v2"
|
||||
print("ok: top-level AG-of-VIDEO encoded per slot")
|
||||
|
||||
# ---- 3. DC-wrapping-AG: branch key untouched, nested AG dict encoded
|
||||
dc_value = {
|
||||
"model": "wan2.7-r2v",
|
||||
"video_refs": {"video1": vid1},
|
||||
"reference_images": {"image1": img1, "image2": img2},
|
||||
"duration": 5,
|
||||
}
|
||||
out = await proxy_node._encode_inputs({"model": dc_value})
|
||||
enc = out["model"]
|
||||
assert enc["model"] == "wan2.7-r2v", enc
|
||||
assert enc["duration"] == 5, enc
|
||||
assert enc["video_refs"]["video1"]["type"] == "video"
|
||||
assert enc["reference_images"]["image1"]["type"] == "image"
|
||||
assert enc["reference_images"]["image2"]["type"] == "image"
|
||||
print("ok: DC-wrapping-AG encoded leaves at every level")
|
||||
|
||||
# ---- 4. AUDIO dict is encoded, NOT recursed into
|
||||
out = await proxy_node._encode_inputs({"audio": audio})
|
||||
assert out["audio"] == {"type": "audio", "encoding": "mp3_base64", "data": "aud"}
|
||||
print("ok: AUDIO dict still routes to encode_audio_input")
|
||||
|
||||
# ---- 5. Already-encoded envelope passes through (no double-encode)
|
||||
pre_encoded = {"type": "image", "encoding": "png_base64", "data": "preimg"}
|
||||
out = await proxy_node._encode_inputs({"images": pre_encoded})
|
||||
assert out["images"] == pre_encoded, out
|
||||
print("ok: pre-encoded envelope passed through unchanged")
|
||||
|
||||
# ---- 6. Plain scalar config dict passes through unchanged
|
||||
out = await proxy_node._encode_inputs(
|
||||
{"config": {"width": 1024, "height": 576, "name": "x"}},
|
||||
)
|
||||
assert out["config"] == {"width": 1024, "height": 576, "name": "x"}, out
|
||||
print("ok: scalar config dict passed through unchanged")
|
||||
|
||||
# ---- 7. _externalize_nested walks per-slot envelopes
|
||||
serialization._upload_log.clear()
|
||||
big_img = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
# encode_image_tensor stub emits ``data="img:1"`` (5 chars). Cap at 3
|
||||
# to force every slot to externalize.
|
||||
out = await proxy_node._encode_inputs(
|
||||
{"reference_images": {"image1": big_img, "image2": big_img}},
|
||||
server_url="https://fake.server",
|
||||
max_inline_bytes=3,
|
||||
)
|
||||
refs = out["reference_images"]
|
||||
assert refs["image1"].get("uri", "").startswith("https://upload/image/"), refs
|
||||
assert refs["image2"].get("uri", "").startswith("https://upload/image/"), refs
|
||||
assert "data" not in refs["image1"] and "data" not in refs["image2"]
|
||||
assert len(serialization._upload_log) == 2, serialization._upload_log
|
||||
print("ok: nested AG envelopes externalized per slot")
|
||||
|
||||
# ---- 7b. _externalize_nested leaves non-envelope siblings alone
|
||||
serialization._upload_log.clear()
|
||||
out = await proxy_node._encode_inputs(
|
||||
{"model": {
|
||||
"model": "wan2.7-r2v",
|
||||
"duration": 5,
|
||||
"reference_images": {"image1": big_img},
|
||||
}},
|
||||
server_url="https://fake.server",
|
||||
max_inline_bytes=3,
|
||||
)
|
||||
enc = out["model"]
|
||||
assert enc["model"] == "wan2.7-r2v" and enc["duration"] == 5, enc
|
||||
assert enc["reference_images"]["image1"]["uri"].startswith("https://upload/image/")
|
||||
print("ok: scalar siblings preserved alongside externalized envelopes")
|
||||
|
||||
# ---- 8. JSON-serializability post-condition
|
||||
payload = await proxy_node._encode_inputs({
|
||||
"reference_images": {"image1": img1, "image2": img2},
|
||||
"video_refs": {"video1": vid1},
|
||||
"audio": audio,
|
||||
"config": {"width": 1024},
|
||||
"scalar": "hello",
|
||||
"n": 42,
|
||||
})
|
||||
json.dumps(payload) # would raise TypeError pre-fix
|
||||
print("ok: full encoded payload is JSON-serializable")
|
||||
|
||||
# ---- 9. Leaf-path regression: single-tensor / single-video / single-audio
|
||||
out = await proxy_node._encode_inputs({"image": img1})
|
||||
assert out["image"]["type"] == "image", out
|
||||
out = await proxy_node._encode_inputs({"video": vid1})
|
||||
assert out["video"]["type"] == "video", out
|
||||
out = await proxy_node._encode_inputs({"audio": audio})
|
||||
assert out["audio"]["type"] == "audio", out
|
||||
out = await proxy_node._encode_inputs({"prompt": "hi", "seed": 7})
|
||||
assert out == {"prompt": "hi", "seed": 7}, out
|
||||
print("ok: top-level leaf path unchanged")
|
||||
|
||||
print("ALL OK")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(asyncio.run(main()))
|
||||
@@ -1,392 +0,0 @@
|
||||
"""Smoke test for the client-side cross-node task_handle decoder.
|
||||
|
||||
Verifies:
|
||||
|
||||
1. ``Capability.IO_TASK_HANDLE`` is intact at the vendored protocol
|
||||
layer and exposed as the well-formed ``io:task_handle`` token.
|
||||
2. ``"task_handle"`` is in ``HEAVY_TYPES`` and ``is_envelope`` returns
|
||||
True for a well-formed task_handle envelope.
|
||||
3. ``decode_task_handle_envelope`` returns a ``TaskHandle`` dataclass
|
||||
with every field round-tripped verbatim (including a non-empty
|
||||
``parent_chain``).
|
||||
4. ``str(handle)`` surfaces the native_id (so UI previews / log lines
|
||||
stay human-meaningful) without leaking the lineage chain.
|
||||
5. ``encode_task_handle(decode_task_handle_envelope(env)) == env`` —
|
||||
the round-trip is lossless so a downstream provider node can re-
|
||||
emit the handle unchanged.
|
||||
6. The top-level ``decode_envelope`` dispatcher routes task_handle
|
||||
envelopes to the new decoder (so proxy_node's generic
|
||||
``_deserialize_output`` picks it up without provider-specific
|
||||
plumbing).
|
||||
7. Malformed task_handle envelopes (missing fields, wrong types,
|
||||
comma-union ``kind``, non-list ``parent_chain``, bad parent_chain
|
||||
entries, wrong encoding) raise ``RnpProtocolError`` with
|
||||
``code=INTERNAL`` — no opaque fallthrough.
|
||||
8. End-to-end against the server's ``make_task_handle_envelope``
|
||||
helper (imported from the comfy-rnp-server worktree) — the
|
||||
envelope the server packs is exactly what the client decoder
|
||||
unpacks, and ``encode_task_handle`` of the decoded result
|
||||
re-produces the same wire shape.
|
||||
9. ``CLIENT_CAPABILITIES`` advertises ``io:task_handle`` so server-
|
||||
side capability-gated descriptors negotiate cleanly.
|
||||
|
||||
Run with the ComfyUI venv (needs torch / comfy_api on PYTHONPATH):
|
||||
|
||||
python notes/run_task_handle_decoder.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Locate the worktree under test.
|
||||
# ---------------------------------------------------------------------------
|
||||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
_CLIENT_DIR = os.path.dirname(_HERE)
|
||||
_COMFYUI_DIR = os.path.dirname(os.path.dirname(_CLIENT_DIR))
|
||||
_SERVER_DIR = os.path.normpath(
|
||||
os.path.join(_COMFYUI_DIR, "..", "comfy-rnp-server")
|
||||
)
|
||||
|
||||
|
||||
def _load_module(name: str, path: str) -> types.ModuleType:
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
# Emulate the client package layout so the relative ``from . import
|
||||
# client as rnp_client`` at the top of serialization.py resolves
|
||||
# without pulling aiohttp / comfy_api_nodes.
|
||||
pkg = types.ModuleType("comfy_remote_nodes_test")
|
||||
pkg.__path__ = [_CLIENT_DIR]
|
||||
sys.modules["comfy_remote_nodes_test"] = pkg
|
||||
|
||||
protocol = _load_module(
|
||||
"comfy_remote_nodes_test.protocol",
|
||||
os.path.join(_CLIENT_DIR, "protocol.py"),
|
||||
)
|
||||
# serialization imports ``client`` for ``rnp_client``; the task_handle
|
||||
# decoder doesn't use it, so a tiny stub suffices.
|
||||
sys.modules["comfy_remote_nodes_test.client"] = types.ModuleType(
|
||||
"comfy_remote_nodes_test.client",
|
||||
)
|
||||
serialization = _load_module(
|
||||
"comfy_remote_nodes_test.serialization",
|
||||
os.path.join(_CLIENT_DIR, "serialization.py"),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Capability constant intact + well-formed token.
|
||||
# ---------------------------------------------------------------------------
|
||||
assert protocol.Capability.IO_TASK_HANDLE == "io:task_handle", (
|
||||
f"Capability.IO_TASK_HANDLE wrong value: "
|
||||
f"{protocol.Capability.IO_TASK_HANDLE!r}"
|
||||
)
|
||||
print("PASS: Capability.IO_TASK_HANDLE = 'io:task_handle'.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. HEAVY_TYPES + is_envelope.
|
||||
# ---------------------------------------------------------------------------
|
||||
assert "task_handle" in protocol.HEAVY_TYPES, (
|
||||
f"'task_handle' missing from HEAVY_TYPES: {sorted(protocol.HEAVY_TYPES)}"
|
||||
)
|
||||
|
||||
hand_envelope = {
|
||||
"type": "task_handle",
|
||||
"encoding": "vendor_inline",
|
||||
"vendor": "tripo",
|
||||
"kind": "MODEL_TASK_ID",
|
||||
"native_id": "abc123-tripo-uuid",
|
||||
"origin_node_id": "42",
|
||||
"parent_chain": [
|
||||
{"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "root-uuid", "origin_node_id": "9"},
|
||||
{"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "prev-uuid", "origin_node_id": "17"},
|
||||
],
|
||||
}
|
||||
|
||||
assert protocol.is_envelope(hand_envelope), (
|
||||
"is_envelope returned False for a well-formed task_handle envelope; "
|
||||
"HEAVY_TYPES update somehow regressed."
|
||||
)
|
||||
print("PASS: 'task_handle' in HEAVY_TYPES and is_envelope recognises it.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. decode_task_handle_envelope returns a populated TaskHandle.
|
||||
# ---------------------------------------------------------------------------
|
||||
handle = serialization.decode_task_handle_envelope(hand_envelope)
|
||||
assert isinstance(handle, serialization.TaskHandle), (
|
||||
f"decode_task_handle_envelope should return TaskHandle, got "
|
||||
f"{type(handle).__name__}"
|
||||
)
|
||||
assert handle.vendor == "tripo"
|
||||
assert handle.kind == "MODEL_TASK_ID"
|
||||
assert handle.native_id == "abc123-tripo-uuid"
|
||||
assert handle.origin_node_id == "42"
|
||||
assert len(handle.parent_chain) == 2
|
||||
assert handle.parent_chain[0] == {
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "root-uuid", "origin_node_id": "9",
|
||||
}
|
||||
assert handle.parent_chain[1] == {
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "prev-uuid", "origin_node_id": "17",
|
||||
}
|
||||
print("PASS: decode_task_handle_envelope returns TaskHandle with full lineage.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. str(handle) surfaces the native_id.
|
||||
# ---------------------------------------------------------------------------
|
||||
assert str(handle) == "abc123-tripo-uuid", (
|
||||
f"str(TaskHandle) should be the native_id; got {str(handle)!r}"
|
||||
)
|
||||
print("PASS: str(TaskHandle) surfaces the native_id.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. encode(decode(env)) == env — lossless round-trip.
|
||||
# ---------------------------------------------------------------------------
|
||||
re_encoded = serialization.encode_task_handle(handle)
|
||||
assert re_encoded == hand_envelope, (
|
||||
f"round-trip mismatch:\n original: {hand_envelope}\n "
|
||||
f"re-encoded: {re_encoded}"
|
||||
)
|
||||
# Empty chain default is preserved.
|
||||
solo = serialization.decode_task_handle_envelope({
|
||||
"type": "task_handle",
|
||||
"encoding": "vendor_inline",
|
||||
"vendor": "kling",
|
||||
"kind": "KLING_VIDEO_ID",
|
||||
"native_id": "video-xyz",
|
||||
"origin_node_id": "5",
|
||||
"parent_chain": [],
|
||||
})
|
||||
assert solo.parent_chain == [], "empty parent_chain should round-trip empty"
|
||||
assert serialization.encode_task_handle(solo)["parent_chain"] == []
|
||||
print("PASS: encode_task_handle round-trips lossless (including empty chain).")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. decode_envelope dispatcher routes to the new decoder.
|
||||
# ---------------------------------------------------------------------------
|
||||
dispatched = serialization.decode_envelope(hand_envelope)
|
||||
assert isinstance(dispatched, serialization.TaskHandle), (
|
||||
f"decode_envelope didn't dispatch task_handle to the right decoder; "
|
||||
f"got {type(dispatched).__name__}"
|
||||
)
|
||||
assert dispatched.native_id == "abc123-tripo-uuid"
|
||||
print("PASS: decode_envelope dispatches 'task_handle' -> decode_task_handle_envelope.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 7. Malformed envelopes raise RnpProtocolError(INTERNAL).
|
||||
# ---------------------------------------------------------------------------
|
||||
def _expect_protocol_error(bad_env: dict, needle: str) -> None:
|
||||
try:
|
||||
serialization.decode_task_handle_envelope(bad_env)
|
||||
except protocol.RnpProtocolError as e:
|
||||
assert e.code == protocol.ErrorCode.INTERNAL, f"wrong error code: {e.code}"
|
||||
assert needle in str(e), f"error msg doesn't mention {needle!r}: {e}"
|
||||
else:
|
||||
raise AssertionError(f"expected RnpProtocolError for {needle!r}")
|
||||
|
||||
|
||||
# Wrong encoding.
|
||||
_expect_protocol_error(
|
||||
{"type": "task_handle", "encoding": "replay_ticket_inline",
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1"},
|
||||
"Unsupported task_handle encoding",
|
||||
)
|
||||
# Missing top-level fields.
|
||||
for missing in ("vendor", "kind", "native_id", "origin_node_id"):
|
||||
bad = {
|
||||
"type": "task_handle", "encoding": "vendor_inline",
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1",
|
||||
}
|
||||
del bad[missing]
|
||||
_expect_protocol_error(bad, repr(missing))
|
||||
# Empty string fields.
|
||||
_expect_protocol_error(
|
||||
{"type": "task_handle", "encoding": "vendor_inline",
|
||||
"vendor": "", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1"},
|
||||
"'vendor'",
|
||||
)
|
||||
# Comma-union kind rejected (input-acceptance string, not emit kind).
|
||||
_expect_protocol_error(
|
||||
{"type": "task_handle", "encoding": "vendor_inline",
|
||||
"vendor": "tripo",
|
||||
"kind": "MODEL_TASK_ID,RIG_TASK_ID,RETARGET_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1"},
|
||||
"singular socket kind",
|
||||
)
|
||||
# parent_chain not a list.
|
||||
_expect_protocol_error(
|
||||
{"type": "task_handle", "encoding": "vendor_inline",
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1",
|
||||
"parent_chain": "not-a-list"},
|
||||
"'parent_chain'",
|
||||
)
|
||||
# Bad parent_chain entry — missing field.
|
||||
_expect_protocol_error(
|
||||
{"type": "task_handle", "encoding": "vendor_inline",
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1",
|
||||
"parent_chain": [{"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "y"}]}, # missing origin_node_id
|
||||
"parent_chain[0]",
|
||||
)
|
||||
# Bad parent_chain entry — comma-union kind in lineage.
|
||||
_expect_protocol_error(
|
||||
{"type": "task_handle", "encoding": "vendor_inline",
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1",
|
||||
"parent_chain": [{"vendor": "tripo",
|
||||
"kind": "MODEL_TASK_ID,RIG_TASK_ID",
|
||||
"native_id": "y", "origin_node_id": "3"}]},
|
||||
"parent_chain[0]",
|
||||
)
|
||||
# Non-dict parent_chain entry.
|
||||
_expect_protocol_error(
|
||||
{"type": "task_handle", "encoding": "vendor_inline",
|
||||
"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "x", "origin_node_id": "1",
|
||||
"parent_chain": ["not-a-dict"]},
|
||||
"parent_chain[0]",
|
||||
)
|
||||
print("PASS: malformed task_handle envelopes raise RnpProtocolError(INTERNAL).")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 8. End-to-end: server make_task_handle_envelope -> client decode.
|
||||
# Skips when the server worktree isn't on disk.
|
||||
# ---------------------------------------------------------------------------
|
||||
server_envelopes_path = os.path.join(
|
||||
_SERVER_DIR, "comfy_rnp_protocol", "envelopes.py",
|
||||
)
|
||||
if os.path.exists(server_envelopes_path):
|
||||
server_pkg = types.ModuleType("comfy_rnp_protocol_e2e_handle")
|
||||
server_pkg.__path__ = [os.path.join(_SERVER_DIR, "comfy_rnp_protocol")]
|
||||
sys.modules["comfy_rnp_protocol_e2e_handle"] = server_pkg
|
||||
_load_module(
|
||||
"comfy_rnp_protocol_e2e_handle.constants",
|
||||
os.path.join(_SERVER_DIR, "comfy_rnp_protocol", "constants.py"),
|
||||
)
|
||||
server_envelopes = _load_module(
|
||||
"comfy_rnp_protocol_e2e_handle.envelopes",
|
||||
os.path.join(_SERVER_DIR, "comfy_rnp_protocol", "envelopes.py"),
|
||||
)
|
||||
assert hasattr(server_envelopes, "make_task_handle_envelope"), (
|
||||
"Server envelopes module missing make_task_handle_envelope — "
|
||||
"server-side PR #39 should have added it."
|
||||
)
|
||||
server_env = server_envelopes.make_task_handle_envelope(
|
||||
vendor="tripo",
|
||||
kind="MODEL_TASK_ID",
|
||||
native_id="server-built-uuid",
|
||||
origin_node_id="100",
|
||||
parent_chain=[
|
||||
{"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": "ancestor-uuid", "origin_node_id": "50"},
|
||||
],
|
||||
)
|
||||
assert server_env["type"] == "task_handle", server_env
|
||||
assert server_env["encoding"] == "vendor_inline", server_env
|
||||
# Client decode of the server-packed bytes.
|
||||
server_decoded = serialization.decode_task_handle_envelope(server_env)
|
||||
assert isinstance(server_decoded, serialization.TaskHandle)
|
||||
assert server_decoded.native_id == "server-built-uuid"
|
||||
assert server_decoded.origin_node_id == "100"
|
||||
assert len(server_decoded.parent_chain) == 1
|
||||
assert server_decoded.parent_chain[0]["native_id"] == "ancestor-uuid"
|
||||
# Round-trip back through encode_task_handle reproduces the server's wire shape.
|
||||
re_encoded_server = serialization.encode_task_handle(server_decoded)
|
||||
assert re_encoded_server == server_env, (
|
||||
f"client encode -> server wire mismatch:\n"
|
||||
f" server: {server_env}\n client: {re_encoded_server}"
|
||||
)
|
||||
# Server-side rejection paths fire as expected (defence in depth —
|
||||
# the client decoder mirrors the same rejections, but a malicious /
|
||||
# buggy producer can't sneak a bad envelope past the server-side
|
||||
# builder either).
|
||||
try:
|
||||
server_envelopes.make_task_handle_envelope(
|
||||
vendor="tripo",
|
||||
kind="MODEL_TASK_ID,RIG_TASK_ID", # comma-union rejected
|
||||
native_id="x",
|
||||
origin_node_id="1",
|
||||
)
|
||||
except ValueError as e:
|
||||
assert "singular socket kind" in str(e), (
|
||||
f"server-side comma-union rejection wrong msg: {e!r}"
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
"server should reject comma-union 'kind' at envelope build time"
|
||||
)
|
||||
# Server enforces parent_chain depth cap (32).
|
||||
try:
|
||||
server_envelopes.make_task_handle_envelope(
|
||||
vendor="tripo",
|
||||
kind="MODEL_TASK_ID",
|
||||
native_id="x",
|
||||
origin_node_id="1",
|
||||
parent_chain=[
|
||||
{"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
"native_id": f"u{i}", "origin_node_id": str(i)}
|
||||
for i in range(33)
|
||||
],
|
||||
)
|
||||
except ValueError as e:
|
||||
assert "depth" in str(e), (
|
||||
f"server-side depth-cap rejection wrong msg: {e!r}"
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
"server should reject parent_chain over MAX_CHAIN_DEPTH"
|
||||
)
|
||||
print(
|
||||
"PASS: end-to-end server make_task_handle_envelope -> client "
|
||||
"decode_task_handle_envelope -> encode_task_handle round-trips "
|
||||
"losslessly, and server-side rejections fire."
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"SKIP end-to-end: server worktree not found at {_SERVER_DIR!r}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 9. CLIENT_CAPABILITIES advertises io:task_handle.
|
||||
# client.py imports aiohttp / comfy_api_nodes — same source-level
|
||||
# contract check we used for the GLB cap in PR #1.
|
||||
# ---------------------------------------------------------------------------
|
||||
client_src = open(
|
||||
os.path.join(_CLIENT_DIR, "client.py"), "r", encoding="utf-8",
|
||||
).read()
|
||||
assert "CLIENT_CAPABILITIES = [" in client_src, (
|
||||
"client.py missing CLIENT_CAPABILITIES list"
|
||||
)
|
||||
cc_start = client_src.index("CLIENT_CAPABILITIES = [")
|
||||
cc_end = client_src.index("]", cc_start)
|
||||
cc_block = client_src[cc_start:cc_end]
|
||||
assert "Capability.IO_TASK_HANDLE" in cc_block, (
|
||||
"Capability.IO_TASK_HANDLE not advertised in CLIENT_CAPABILITIES"
|
||||
)
|
||||
print("PASS: CLIENT_CAPABILITIES advertises Capability.IO_TASK_HANDLE.")
|
||||
|
||||
print("ALL CHECKS PASSED.")
|
||||
+1
-53
@@ -225,61 +225,9 @@ class Capability:
|
||||
# ``uri`` per file for URL-fetch-out-of-band semantics; not in
|
||||
# this PR.
|
||||
MODEL_3D_BUNDLE_INLINE = "model_3d:bundle_inline"
|
||||
# Cross-node task handle envelope: wraps a vendor-native task /
|
||||
# video / job ID (Tripo MODEL_TASK_ID, Kling video_id, future
|
||||
# partner handles) with enough lineage metadata to chain provider
|
||||
# nodes (Tripo Texture/Refine/Rig/Retarget/Conversion, Kling
|
||||
# VideoExtend, ...) across a single workflow graph. Wire shape:
|
||||
#
|
||||
# {
|
||||
# "type": "task_handle",
|
||||
# "encoding": "vendor_inline",
|
||||
# "vendor": "tripo",
|
||||
# "kind": "MODEL_TASK_ID",
|
||||
# "native_id": "abc123-tripo-uuid",
|
||||
# "origin_node_id": "42",
|
||||
# "parent_chain": [
|
||||
# {"vendor": "tripo", "kind": "MODEL_TASK_ID",
|
||||
# "native_id": "prev-uuid", "origin_node_id": "17"}
|
||||
# ]
|
||||
# }
|
||||
#
|
||||
# ``vendor`` is a lowercase provider slug. ``kind`` is the
|
||||
# singular emitted socket kind string (``MODEL_TASK_ID`` /
|
||||
# ``RIG_TASK_ID`` / ``RETARGET_TASK_ID`` / ``KLING_VIDEO_ID``) —
|
||||
# NEVER the comma-union string that provider nodes use for input
|
||||
# type-acceptance (e.g. ``TripoConversionNode`` accepts
|
||||
# ``"MODEL_TASK_ID,RIG_TASK_ID,RETARGET_TASK_ID"``; the wire kind
|
||||
# is whatever the upstream actually emitted). ``native_id`` is
|
||||
# the vendor-native opaque string. ``origin_node_id`` is the
|
||||
# upstream workflow node's unique id (log correlation + future
|
||||
# replay).
|
||||
#
|
||||
# ``parent_chain`` is a FLAT list (not recursive nested
|
||||
# envelopes) of ``{vendor, kind, native_id, origin_node_id}``
|
||||
# refs in root-to-immediate-parent order. Each downstream node
|
||||
# appends its parent's ref before the parent's chain is copied
|
||||
# through; the chain captures the full lineage. Hard-capped at
|
||||
# depth 32 server-side to prevent runaway growth.
|
||||
#
|
||||
# Negotiation: server emits ``task_handle`` envelopes only when
|
||||
# the inbound request advertises this capability; legacy clients
|
||||
# surface NEGOTIATION_FAILED at descriptor-load time on any
|
||||
# provider that publishes a handle-typed output OR accepts one
|
||||
# as input. Whole-descriptor gating — partial-socket synthesis
|
||||
# is rejected because the chain breaks meaningfully if either
|
||||
# end is missing.
|
||||
#
|
||||
# The client decoder returns a ``TaskHandle`` dataclass (NOT a
|
||||
# plain string) so lineage metadata round-trips when the value
|
||||
# is passed through to the next provider node — re-encoding the
|
||||
# same dataclass on a downstream submit preserves ``parent_chain``
|
||||
# without the proxy_node needing to know anything provider-
|
||||
# specific.
|
||||
IO_TASK_HANDLE = "io:task_handle"
|
||||
|
||||
|
||||
HEAVY_TYPES = frozenset({"image", "video", "audio", "mask", "model_3d", "task_handle"})
|
||||
HEAVY_TYPES = frozenset({"image", "video", "audio", "mask", "model_3d"})
|
||||
|
||||
|
||||
def is_envelope(value: Any) -> bool:
|
||||
|
||||
+101
-7
@@ -695,6 +695,30 @@ def _build_input(name: str, io_type: Any, options: dict[str, Any], optional: boo
|
||||
return IO.Audio.Input(name, **primitive_common)
|
||||
if io_type == "MASK":
|
||||
return IO.Mask.Input(name, **primitive_common)
|
||||
if io_type == "MODEL_3D":
|
||||
# 3D MODEL input: symmetric counterpart of the ``"MODEL_3D"``
|
||||
# entry in ``_OUTPUT_CLASSES`` above. The Tencent Hunyuan3D
|
||||
# MODEL_3D-input nodes (``TencentModelTo3DUVNode`` /
|
||||
# ``Tencent3DTextureEditNode`` / ``Tencent3DPartNode`` /
|
||||
# ``TencentSmartTopologyNode``) all declare their model_3d
|
||||
# input as ``IO.MultiType.Input(types=[IO.File3DGLB,
|
||||
# IO.File3DOBJ, IO.File3DFBX, IO.File3DAny])`` so the socket
|
||||
# accepts any upstream 3D source: ``Load3D``'s ``File3DAny``
|
||||
# output, Meshy / Hunyuan3D / Tripo / Rodin's per-format
|
||||
# outputs (``FILE_3D_GLB`` / ``FILE_3D_OBJ`` / ``FILE_3D_FBX``),
|
||||
# etc. Mirror that union here so the proxy socket accepts the
|
||||
# same set of comfytypes upstream does. Per-node format
|
||||
# validation (e.g. 3DTextureEdit's FBX-only enforcement) is
|
||||
# done server-side in the provider's ``execute()`` against
|
||||
# the inline envelope's ``format`` extra; this socket-level
|
||||
# filter is purely about frontend connection validity.
|
||||
return IO.MultiType.Input(
|
||||
name,
|
||||
types=[
|
||||
IO.File3DGLB, IO.File3DOBJ, IO.File3DFBX, IO.File3DAny,
|
||||
],
|
||||
**primitive_common,
|
||||
)
|
||||
if io_type == "DYNAMIC_COMBO":
|
||||
# Wire shape: ``["DYNAMIC_COMBO", {"options": [{"key": str,
|
||||
# "inputs": [[name, io_type, opts], ...]}, ...], "tooltip": ...}]``.
|
||||
@@ -1546,6 +1570,8 @@ async def _encode_inputs(
|
||||
* ``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.
|
||||
* ``VideoInput`` subclass (duck-typed via ``save_to`` +
|
||||
``get_duration``) → mp4-base64 video envelope.
|
||||
* Everything else passes through as-is (already JSON-serializable).
|
||||
|
||||
When ``max_inline_bytes`` is set, any envelope whose inline payload
|
||||
@@ -1574,17 +1600,65 @@ async def _encode_inputs(
|
||||
name, value, validate_map.get(name) or {}, node_id=node_id,
|
||||
)
|
||||
encoded = _encode_one(name, value, serialization_map.get(name))
|
||||
# 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,
|
||||
# Externalize any envelope payloads, walking into nested dicts so
|
||||
# AUTOGROW slot values and DYNAMIC_COMBO sub-inputs whose encoded
|
||||
# form is a per-slot envelope are uploaded the same way a
|
||||
# top-level singleton envelope is. ``_externalize_nested`` is a
|
||||
# no-op on plain scalars and on dicts that contain no envelopes.
|
||||
encoded = await _externalize_nested(
|
||||
encoded,
|
||||
server_url=server_url,
|
||||
max_inline_bytes=max_inline_bytes,
|
||||
auth_headers=auth_headers,
|
||||
)
|
||||
out[name] = encoded
|
||||
return out
|
||||
|
||||
|
||||
async def _externalize_nested(
|
||||
value: Any,
|
||||
*,
|
||||
server_url: str | None,
|
||||
max_inline_bytes: int | None,
|
||||
auth_headers: dict[str, str] | None,
|
||||
) -> Any:
|
||||
"""Recursively externalize any RNP envelopes inside ``value``.
|
||||
|
||||
Mirrors :func:`_encode_one`'s dict-recursion: ComfyUI hands the
|
||||
proxy_node nested AUTOGROW / DYNAMIC_COMBO runtime dicts
|
||||
(e.g. ``{"reference_images": {"image1": <env>, "image2": <env>}}``
|
||||
or ``{"model": {"branch": "wan2.7-r2v",
|
||||
"video_refs": {"video1": <env>}}}``); without this walk the
|
||||
per-slot envelopes would ship inline regardless of
|
||||
``max_inline_bytes`` because :func:`_encode_inputs` only inspected
|
||||
the top-level value.
|
||||
|
||||
Stops at envelope leaves (delegates the per-envelope decision to
|
||||
:func:`serialization.maybe_externalize`) and at plain scalars (leaves
|
||||
them unchanged). Lists are intentionally left alone — AG/DC runtime
|
||||
shapes always materialize as dicts, never lists, and recursing into
|
||||
arbitrary lists would touch opaque payloads. Recursion is bounded
|
||||
by the AUTOGROW max slot count (100) and the maximum DC branch
|
||||
nesting depth (1 today), so an in-place ``async def`` walk is fine.
|
||||
"""
|
||||
if serialization.is_envelope(value):
|
||||
return await serialization.maybe_externalize(
|
||||
value,
|
||||
server_url=server_url,
|
||||
max_inline_bytes=max_inline_bytes,
|
||||
auth_headers=auth_headers,
|
||||
)
|
||||
if isinstance(value, dict):
|
||||
out: dict[str, Any] = {}
|
||||
for k, v in value.items():
|
||||
out[k] = await _externalize_nested(
|
||||
v,
|
||||
server_url=server_url,
|
||||
max_inline_bytes=max_inline_bytes,
|
||||
auth_headers=auth_headers,
|
||||
)
|
||||
out[name] = encoded
|
||||
return out
|
||||
return out
|
||||
return value
|
||||
|
||||
|
||||
def _enforce_local_validate(
|
||||
@@ -1692,8 +1766,28 @@ def _encode_one(
|
||||
value: Any,
|
||||
accepted_encodings: list[str] | None = None,
|
||||
) -> Any:
|
||||
# ComfyUI's V3 ``build_nested_inputs`` rebuilds AUTOGROW / DYNAMIC_COMBO
|
||||
# runtime kwargs into nested dicts (``{"reference_images": {"image1":
|
||||
# <tensor>, ...}}`` for top-level AG; ``{"<dc-name>": branch_key,
|
||||
# "<nested-ag>": {"video1": <VideoInput>, ...}}`` for AG-inside-DC).
|
||||
# Recurse into plain dicts so heavy leaves anywhere in that tree get
|
||||
# encoded — without this an AG-of-IMAGE slot would ship the raw
|
||||
# torch.Tensor inside a dict and crash at ``json.dumps`` time. Stops
|
||||
# at envelope dicts (already-encoded payloads) so we don't double-
|
||||
# encode wire-shape values that happen to be dicts.
|
||||
if serialization.is_envelope(value):
|
||||
return value
|
||||
if isinstance(value, dict) and not serialization.is_audio_input(value):
|
||||
return {
|
||||
k: _encode_one(f"{name}.{k}", v, accepted_encodings)
|
||||
for k, v in value.items()
|
||||
}
|
||||
if serialization.is_audio_input(value):
|
||||
return serialization.encode_audio_input(value)
|
||||
if serialization.is_video_input(value):
|
||||
return serialization.encode_video_input(value)
|
||||
if serialization.is_model3d_input(value):
|
||||
return serialization.encode_model3d_input(value)
|
||||
if not serialization._is_torch_tensor(value):
|
||||
return value
|
||||
rank = value.dim()
|
||||
|
||||
+147
-159
@@ -465,9 +465,44 @@ def decode_audio_envelope(envelope: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Video decode (encode lands when a remote node accepts VIDEO inputs)
|
||||
# Video encode / decode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def encode_video_input(video: Any) -> dict[str, Any]:
|
||||
"""Encode a ComfyUI VIDEO (``VideoInput`` subclass) as an mp4-base64
|
||||
video envelope.
|
||||
|
||||
Mirrors the AUDIO encoder above: re-encodes the Video object to an
|
||||
in-memory mp4/H.264 byte buffer (matches upstream
|
||||
``comfy_api_nodes/util/conversions.py`` ``video_to_base64_string``)
|
||||
and base64-encodes it. Populates ``duration_s`` metadata when the
|
||||
Video object exposes ``get_duration()`` so server-side providers
|
||||
(e.g. ``GrokVideoExtendProvider._validate_video_duration_envelope``)
|
||||
can range-check without demuxing the MP4.
|
||||
"""
|
||||
from comfy_api_nodes.util.conversions import video_to_base64_string
|
||||
|
||||
b64 = video_to_base64_string(video)
|
||||
extra: dict[str, Any] = {}
|
||||
duration_s: float | None = None
|
||||
try:
|
||||
getter = getattr(video, "get_duration", None)
|
||||
if callable(getter):
|
||||
d = getter()
|
||||
if d is not None:
|
||||
duration_s = float(d)
|
||||
except Exception:
|
||||
duration_s = None
|
||||
if duration_s is not None:
|
||||
extra["duration_s"] = duration_s
|
||||
return {
|
||||
"type": "video",
|
||||
"encoding": "mp4_base64",
|
||||
"data": b64,
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def decode_video_envelope(envelope: dict[str, Any]) -> Any:
|
||||
"""Decode a video envelope into a ComfyUI Video object (mp4 inline)."""
|
||||
encoding = envelope.get("encoding")
|
||||
@@ -481,9 +516,69 @@ def decode_video_envelope(envelope: dict[str, Any]) -> Any:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3D model decode (encode lands when a remote node accepts MODEL_3D inputs)
|
||||
# 3D model encode / decode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def encode_model3d_input(model_3d: Any) -> dict[str, Any]:
|
||||
"""Encode a ComfyUI ``File3D`` (or compatible) as an inline
|
||||
single-file 3D-model envelope.
|
||||
|
||||
Mirrors the AUDIO / VIDEO encoders above: reads the primary mesh
|
||||
bytes (via the ``File3D`` API: ``get_data().read()`` returns a
|
||||
``BytesIO`` payload; falling back to ``read_bytes()`` for older
|
||||
sources) and base64-encodes them. The envelope's ``encoding`` is
|
||||
``"glb_inline"`` (the generic single-file inline encoding the
|
||||
server-side decoder dispatches on at :func:`decode_model3d_envelope`
|
||||
— the decoder hands the bytes to ``File3D(BytesIO(...), format)``
|
||||
regardless of whether the actual format is glb / obj / fbx / etc.,
|
||||
so the ``encoding`` name is a logical "inline single-file 3D" tag
|
||||
and the ``format`` extra picks the file extension).
|
||||
|
||||
Multi-file bundles (``BundledFile3D``-style values where every
|
||||
companion file matters — e.g. an OBJ with a sibling .mtl + texture
|
||||
PNGs) are not yet supported on the encode path: this PR ships the
|
||||
minimal wire shape needed for the Tencent Hunyuan3D MODEL_3D-input
|
||||
nodes which all upload a single file via ``upload_3d_model_to_comfyapi``
|
||||
upstream. Bundle encode (``bundle_inline``) lands when a partner
|
||||
node actually needs the companion bytes.
|
||||
"""
|
||||
file_format = ""
|
||||
fmt_attr = getattr(model_3d, "format", None)
|
||||
if isinstance(fmt_attr, str):
|
||||
file_format = fmt_attr.lower()
|
||||
|
||||
# ``File3D.get_data()`` returns a ``BytesIO`` (per
|
||||
# ``comfy_api.latest._util.geometry_types.File3D``). Other File3D-
|
||||
# compatible classes may instead expose ``read_bytes`` returning
|
||||
# raw bytes; fall back to that for compatibility.
|
||||
raw: bytes
|
||||
get_data = getattr(model_3d, "get_data", None)
|
||||
if callable(get_data):
|
||||
buf = get_data()
|
||||
if hasattr(buf, "seek"):
|
||||
buf.seek(0)
|
||||
raw = buf.read() if hasattr(buf, "read") else bytes(buf)
|
||||
else:
|
||||
read_bytes = getattr(model_3d, "read_bytes", None)
|
||||
if callable(read_bytes):
|
||||
raw = read_bytes()
|
||||
else:
|
||||
raise RnpProtocolError(
|
||||
f"unsupported 3D model value: {type(model_3d).__name__!r} "
|
||||
f"has no get_data() or read_bytes()",
|
||||
code=ErrorCode.INTERNAL,
|
||||
)
|
||||
|
||||
envelope: dict[str, Any] = {
|
||||
"type": "model_3d",
|
||||
"encoding": "glb_inline",
|
||||
"data": base64.b64encode(raw).decode("ascii"),
|
||||
}
|
||||
if file_format:
|
||||
envelope["format"] = file_format
|
||||
return envelope
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multi-file 3D bundle: a ``File3D``-compatible class that carries the
|
||||
# primary mesh plus N companion files (textures / .mtl / .bin / etc.).
|
||||
@@ -795,165 +890,16 @@ def decode_model3d_envelope(envelope: dict[str, Any]) -> Any:
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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,
|
||||
"image": decode_image_envelope,
|
||||
"mask": decode_mask_envelope,
|
||||
"audio": decode_audio_envelope,
|
||||
"video": decode_video_envelope,
|
||||
"model_3d": decode_model3d_envelope,
|
||||
}
|
||||
|
||||
|
||||
@@ -986,6 +932,47 @@ def is_audio_input(value: Any) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def is_video_input(value: Any) -> bool:
|
||||
"""Best-effort check for a ComfyUI ``VideoInput`` subclass.
|
||||
|
||||
Duck-types on the ``save_to`` + ``get_duration`` method pair: both
|
||||
are defined on the ``VideoInput`` abstract base in
|
||||
``comfy_api.latest._input`` and present on every concrete subclass
|
||||
(``VideoFromFile`` / ``VideoFromComponents``). Avoids importing
|
||||
``VideoInput`` at module load (which would pull torch via the
|
||||
``comfy_api`` tree).
|
||||
"""
|
||||
if isinstance(value, dict):
|
||||
return False
|
||||
return (
|
||||
callable(getattr(value, "save_to", None))
|
||||
and callable(getattr(value, "get_duration", None))
|
||||
)
|
||||
|
||||
|
||||
def is_model3d_input(value: Any) -> bool:
|
||||
"""Best-effort check for a ComfyUI ``File3D`` (or compatible).
|
||||
|
||||
Duck-types on the ``format`` string attribute + ``get_data``
|
||||
callable: both are present on the ``File3D`` class in
|
||||
``comfy_api.latest._util.geometry_types`` and the
|
||||
``BundledFile3D`` subclass built lazily in
|
||||
:func:`_bundled_file3d_class`. Avoids importing ``File3D`` at
|
||||
module load (which would pull torch via the ``comfy_api`` tree).
|
||||
|
||||
The ``VideoInput`` shape (``save_to`` + ``get_duration``) does not
|
||||
expose a ``format`` string property, so VIDEO and MODEL_3D values
|
||||
don't collide; caller (:func:`_encode_one`) still dispatches
|
||||
AUDIO / VIDEO before MODEL_3D as defense-in-depth.
|
||||
"""
|
||||
if isinstance(value, dict):
|
||||
return False
|
||||
return (
|
||||
isinstance(getattr(value, "format", None), str)
|
||||
and callable(getattr(value, "get_data", None))
|
||||
)
|
||||
|
||||
|
||||
def _is_torch_tensor(value: Any) -> bool:
|
||||
try:
|
||||
import torch
|
||||
@@ -1002,15 +989,16 @@ __all__ = [
|
||||
"decode_mask_envelope",
|
||||
"encode_audio_input",
|
||||
"decode_audio_envelope",
|
||||
"encode_video_input",
|
||||
"decode_video_envelope",
|
||||
"encode_model3d_input",
|
||||
"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",
|
||||
"is_video_input",
|
||||
"is_model3d_input",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user