diff --git a/client.py b/client.py index f256ad1..35efb8e 100644 --- a/client.py +++ b/client.py @@ -91,6 +91,18 @@ CLIENT_CAPABILITIES = [ # legacy client without the decoder surfaces NEGOTIATION_FAILED # at descriptor-load time rather than crashing at execute. Capability.MODEL_3D_GLB_INLINE, + # 3D MODEL BUNDLE envelope (multi-file: primary mesh plus N + # companion files — textures / .mtl / .bin / etc. — as a single + # coherent unit). The decoder in + # ``serialization.decode_model3d_envelope`` dispatches on + # ``encoding="bundle_inline"`` and returns a ``BundledFile3D`` + # (a ``File3D`` subclass) whose ``get_source()`` / ``save_to()`` + # materialise the whole bundle on disk so cross-file refs (.obj + # -> .mtl -> textures; .gltf -> .bin + textures) resolve via + # the filesystem. Servers gate bundle-emitting descriptors on + # this token; legacy clients surface NEGOTIATION_FAILED at + # descriptor-load time rather than crashing at execute. + Capability.MODEL_3D_BUNDLE_INLINE, ] # Reported in the X-RNP-Client-Version header. diff --git a/notes/run_model3d_bundle_decoder.py b/notes/run_model3d_bundle_decoder.py new file mode 100644 index 0000000..9c820ee --- /dev/null +++ b/notes/run_model3d_bundle_decoder.py @@ -0,0 +1,461 @@ +"""Smoke test for the client-side multi-file MODEL_3D bundle decoder. + +Verifies: + +1. ``Capability.MODEL_3D_BUNDLE_INLINE`` is intact at the vendored + protocol layer and exposed as the well-formed + ``model_3d:bundle_inline`` token. +2. ``is_envelope`` recognises a ``bundle_inline`` envelope (still + ``type=model_3d`` — HEAVY_TYPES unchanged). +3. ``decode_model3d_envelope`` dispatches ``encoding=bundle_inline`` + to the new ``BundledFile3D`` class, returns a real ``File3D`` + subclass, and round-trips the primary mesh bytes via + ``get_data()`` / ``get_bytes()`` without materialising anything + on disk. +4. ``get_source()`` materialises the whole bundle into a temp + directory using the producer-supplied relative paths, so a .obj + that references its .mtl by filename (and a texture in + ``textures/``) finds them on disk via the filesystem. +5. ``save_to(dest)`` copies the whole bundle into the destination + directory — the primary file under the caller-specified name, + sidecars under their ORIGINAL relative paths — so cross-file refs + inside the .obj (which point at ``model.mtl``) still resolve when + the user saves the model somewhere stable. +6. Bundle-specific accessors (``primary_path``, ``file_paths``, + ``file_bytes``, ``file_role``, ``paths_with_role``) return the + expected values; unknown roles round-trip without error. +7. Malformed bundle envelopes (missing primary_path, missing + files, missing data, primary_path not in files) raise + ``RnpProtocolError`` with ``code=INTERNAL`` — no opaque + fallthrough. +8. End-to-end against the server's ``make_model3d_bundle_envelope`` + helper (imported from the comfy-rnp-server worktree) — the bytes + the server packs into each file are exactly the bytes the + client unpacks, and the path sanitization happens server-side so + the client never sees a malicious path. +9. ``CLIENT_CAPABILITIES`` advertises ``model_3d:bundle_inline`` + so server-side capability-gated descriptors negotiate cleanly. + +Run with the ComfyUI venv (needs torch / comfy_api on PYTHONPATH): + + python notes/run_model3d_bundle_decoder.py +""" +from __future__ import annotations + +import base64 +import importlib.util +import os +import shutil +import sys +import tempfile +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 + + +# Make ComfyUI importable so ``comfy_api.latest._util.geometry_types.File3D`` +# resolves at first BundledFile3D construction. +if _COMFYUI_DIR not in sys.path: + sys.path.insert(0, _COMFYUI_DIR) +if _SERVER_DIR not in sys.path: + sys.path.insert(0, _SERVER_DIR) + +# ComfyUI's full ``comfy_api`` tree pulls torch / PIL / comfy_execution at +# import time. Stub just what BundledFile3D needs — the leaf +# ``comfy_api.latest._util.geometry_types.File3D`` class with a faithful +# subclass-friendly implementation. +class _StubFile3D: + """Faithful subset of upstream ``comfy_api.latest._util.geometry_types.File3D``. + + Implements the same parent-class shape (``_source`` / ``_format`` + attrs + ``get_source`` / ``get_data`` / ``get_bytes`` / ``save_to`` + methods) so subclassing it produces an object indistinguishable + from the real thing for the decoder's purposes. The smoke test + deliberately stubs torch + the heavy ``comfy_api.latest._util`` + package init to keep the test runnable in the rnp venv. + """ + def __init__(self, source, file_format: str = "") -> None: + self._source = source + self._format = (file_format or "").lstrip(".").lower() + + @property + def format(self) -> str: + return self._format + + @format.setter + def format(self, value: str) -> None: + self._format = (value or "").lstrip(".").lower() + + +# Stub ``comfy_api.latest._util`` and the nested ``geometry_types`` module +# so ``from comfy_api.latest._util.geometry_types import File3D`` works +# without pulling torch / PIL / etc. +_comfy_api = types.ModuleType("comfy_api") +_comfy_api.__path__ = [] # type: ignore[attr-defined] +_comfy_api_latest = types.ModuleType("comfy_api.latest") +_comfy_api_latest.__path__ = [] # type: ignore[attr-defined] +_comfy_api_latest_util = types.ModuleType("comfy_api.latest._util") +_comfy_api_latest_util.__path__ = [] # type: ignore[attr-defined] +_comfy_api_latest_util.File3D = _StubFile3D +_comfy_api_latest_util_geom = types.ModuleType("comfy_api.latest._util.geometry_types") +_comfy_api_latest_util_geom.File3D = _StubFile3D +sys.modules.setdefault("comfy_api", _comfy_api) +sys.modules.setdefault("comfy_api.latest", _comfy_api_latest) +sys.modules.setdefault("comfy_api.latest._util", _comfy_api_latest_util) +sys.modules.setdefault("comfy_api.latest._util.geometry_types", _comfy_api_latest_util_geom) + +# Emulate the client package layout for the relative imports. +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 bundle 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"), +) + + +# --------------------------------------------------------------------------- +# Build a minimal Wavefront-OBJ + MTL + PNG-stub bundle for the round-trip. +# These bytes don't need to be loadable by an actual OBJ parser — the +# decoder is bytes-in / bytes-out — but using realistic content makes +# debugging easier if a future regression breaks the cross-file refs. +# --------------------------------------------------------------------------- +def _make_minimal_obj() -> bytes: + return ( + b"# Minimal OBJ for RNP bundle smoke test.\n" + b"mtllib model.mtl\n" + b"v 0.0 0.0 0.0\nv 1.0 0.0 0.0\nv 0.0 1.0 0.0\n" + b"usemtl test_material\n" + b"f 1 2 3\n" + ) + + +def _make_minimal_mtl() -> bytes: + return ( + b"# Minimal MTL for RNP bundle smoke test.\n" + b"newmtl test_material\n" + b"Ka 1.0 1.0 1.0\n" + b"Kd 1.0 1.0 1.0\n" + b"map_Kd textures/albedo.png\n" + b"map_bump textures/normal.png\n" + ) + + +def _make_stub_png(tag: bytes) -> bytes: + # Not a real PNG — the decoder doesn't parse it. Tag distinguishes + # each "texture" so the round-trip assertions can pin which bytes + # went where. + return b"\x89PNG\r\n\x1a\n" + tag + + +obj_bytes = _make_minimal_obj() +mtl_bytes = _make_minimal_mtl() +albedo_bytes = _make_stub_png(b"ALBEDO_STUB") +normal_bytes = _make_stub_png(b"NORMAL_STUB") +roughness_bytes = _make_stub_png(b"ROUGHNESS_STUB") + + +# --------------------------------------------------------------------------- +# 1. Capability constant intact at the vendored protocol layer. +# --------------------------------------------------------------------------- +assert protocol.Capability.MODEL_3D_BUNDLE_INLINE == "model_3d:bundle_inline", ( + f"Capability.MODEL_3D_BUNDLE_INLINE wrong value: " + f"{protocol.Capability.MODEL_3D_BUNDLE_INLINE!r}" +) +print("PASS: Capability.MODEL_3D_BUNDLE_INLINE = 'model_3d:bundle_inline'.") + + +# --------------------------------------------------------------------------- +# Build a hand-crafted bundle envelope for the next several checks. +# --------------------------------------------------------------------------- +hand_envelope = { + "type": "model_3d", + "encoding": "bundle_inline", + "format": "obj", + "primary_path": "model.obj", + "files": [ + {"path": "model.obj", "format": "obj", "role": "mesh", + "data": base64.b64encode(obj_bytes).decode("ascii"), + "byte_size": len(obj_bytes)}, + {"path": "model.mtl", "format": "mtl", "role": "material", + "data": base64.b64encode(mtl_bytes).decode("ascii")}, + {"path": "textures/albedo.png", "format": "png", "role": "texture_diffuse", + "data": base64.b64encode(albedo_bytes).decode("ascii")}, + {"path": "textures/normal.png", "format": "png", "role": "texture_normal", + "data": base64.b64encode(normal_bytes).decode("ascii")}, + {"path": "textures/roughness.png","format": "png","role": "texture_roughness", + "data": base64.b64encode(roughness_bytes).decode("ascii")}, + # Unknown role — must round-trip without error. + {"path": "extras/notes.txt", "role": "producer_notes", + "data": base64.b64encode(b"hello").decode("ascii")}, + ], +} + + +# --------------------------------------------------------------------------- +# 2. is_envelope recognises a bundle envelope. +# --------------------------------------------------------------------------- +assert protocol.is_envelope(hand_envelope), ( + "is_envelope returned False for a well-formed bundle_inline envelope; " + "HEAVY_TYPES update somehow regressed." +) +print("PASS: is_envelope recognises bundle_inline envelopes.") + + +# --------------------------------------------------------------------------- +# 3. decode_model3d_envelope dispatches + round-trips primary bytes. +# --------------------------------------------------------------------------- +bundle = serialization.decode_model3d_envelope(hand_envelope) +bundle_cls = serialization._bundled_file3d_class() +assert isinstance(bundle, bundle_cls), ( + f"decode_model3d_envelope should return BundledFile3D, got " + f"{type(bundle).__name__}" +) +# It must ALSO be a File3D (parent class) — that's the whole point of +# the subclass approach. +assert isinstance(bundle, _StubFile3D), ( + f"BundledFile3D must subclass File3D — got {type(bundle).__mro__}" +) +# Primary file bytes round-trip via get_data / get_bytes (no disk I/O). +assert bundle.get_bytes() == obj_bytes, ( + f"primary OBJ bytes lost: {len(bundle.get_bytes())} vs {len(obj_bytes)}" +) +src = bundle.get_data() +src.seek(0) +assert src.read() == obj_bytes, "primary OBJ get_data() didn't round-trip" +assert bundle.format == "obj", f"primary format wrong: {bundle.format!r}" +print("PASS: decode_model3d_envelope returns BundledFile3D + primary bytes round-trip.") + + +# --------------------------------------------------------------------------- +# 4. get_source() materialises the whole bundle on disk. +# --------------------------------------------------------------------------- +disk_path = bundle.get_source() +assert os.path.isfile(disk_path), f"get_source() should return existing path: {disk_path}" +assert disk_path.endswith("model.obj"), f"primary path wrong: {disk_path}" +with open(disk_path, "rb") as f: + assert f.read() == obj_bytes, "disk-materialised primary OBJ bytes mismatched" +# Sidecars must be at the SAME relative paths so the .obj's mtllib / +# map_Kd refs resolve. +temp_root = bundle.materialize_to_temp_dir() +assert os.path.isfile(os.path.join(temp_root, "model.mtl")), ( + "model.mtl missing from materialised bundle" +) +assert os.path.isfile(os.path.join(temp_root, "textures", "albedo.png")), ( + "textures/albedo.png missing from materialised bundle" +) +assert os.path.isfile(os.path.join(temp_root, "textures", "normal.png")) +assert os.path.isfile(os.path.join(temp_root, "textures", "roughness.png")) +assert os.path.isfile(os.path.join(temp_root, "extras", "notes.txt")) +# Sidecar bytes match. +with open(os.path.join(temp_root, "model.mtl"), "rb") as f: + assert f.read() == mtl_bytes +with open(os.path.join(temp_root, "textures", "albedo.png"), "rb") as f: + assert f.read() == albedo_bytes +print("PASS: get_source() materialises whole bundle on disk at producer-supplied paths.") + + +# --------------------------------------------------------------------------- +# 5. save_to(dest) copies the whole bundle into the destination directory. +# --------------------------------------------------------------------------- +save_dir = tempfile.mkdtemp(prefix="rnp_save_to_test_") +try: + out_path = bundle.save_to(os.path.join(save_dir, "renamed_scene.obj")) + assert os.path.isfile(out_path) + with open(out_path, "rb") as f: + assert f.read() == obj_bytes, "primary OBJ bytes wrong after save_to" + # Sidecars keep ORIGINAL filenames so cross-file refs inside the + # OBJ (which point at "model.mtl") still resolve relative to the + # renamed primary file's directory. + assert os.path.isfile(os.path.join(save_dir, "model.mtl")), ( + "save_to should copy model.mtl alongside the renamed primary" + ) + assert os.path.isfile(os.path.join(save_dir, "textures", "albedo.png")) + assert os.path.isfile(os.path.join(save_dir, "textures", "normal.png")) + with open(os.path.join(save_dir, "textures", "normal.png"), "rb") as f: + assert f.read() == normal_bytes + print("PASS: save_to copies whole bundle, primary renamed, sidecars preserve filenames.") +finally: + shutil.rmtree(save_dir, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# 6. Bundle-specific accessors. +# --------------------------------------------------------------------------- +assert bundle.primary_path == "model.obj" +assert set(bundle.file_paths) == { + "model.obj", "model.mtl", + "textures/albedo.png", "textures/normal.png", "textures/roughness.png", + "extras/notes.txt", +} +assert bundle.file_bytes("textures/albedo.png") == albedo_bytes +# Optional roles. +assert bundle.file_role("model.obj") == "mesh" +assert bundle.file_role("textures/normal.png") == "texture_normal" +# Unknown roles round-trip — no enum check. +assert bundle.file_role("extras/notes.txt") == "producer_notes" +# Missing role (none supplied) returns None. +assert bundle.file_role("nonexistent.png") is None +# Convenience lookup. +assert bundle.paths_with_role("texture_normal") == ["textures/normal.png"] +assert bundle.paths_with_role("nonexistent_role") == [] +print("PASS: bundle accessors (primary_path / file_paths / file_role / paths_with_role).") + + +# --------------------------------------------------------------------------- +# 7. Malformed bundle envelopes raise RnpProtocolError with code=INTERNAL. +# --------------------------------------------------------------------------- +def _expect_protocol_error(bad_env: dict, needle: str) -> None: + try: + serialization.decode_model3d_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}") + +_expect_protocol_error( + {"type": "model_3d", "encoding": "bundle_inline", "files": [ + {"path": "x.obj", "data": "AA=="}, + ]}, + "primary_path", +) +_expect_protocol_error( + {"type": "model_3d", "encoding": "bundle_inline", "primary_path": "x.obj"}, + "files", +) +_expect_protocol_error( + {"type": "model_3d", "encoding": "bundle_inline", "primary_path": "x.obj", + "files": [{"path": "y.obj", "data": "AA=="}]}, + "primary_path 'x.obj' not in", +) +_expect_protocol_error( + {"type": "model_3d", "encoding": "bundle_inline", "primary_path": "x.obj", + "files": [{"path": "x.obj"}]}, # missing data + "missing inline 'data'", +) +_expect_protocol_error( + {"type": "model_3d", "encoding": "wat_inline", "data": "AA=="}, + "Unsupported model_3d encoding", +) +print("PASS: malformed bundle envelopes raise RnpProtocolError(INTERNAL).") + + +# --------------------------------------------------------------------------- +# 8. End-to-end: server make_model3d_bundle_envelope -> client decode. +# Skips when the server worktree isn't on disk. +# --------------------------------------------------------------------------- +server_protocol_path = os.path.join( + _SERVER_DIR, "comfy_rnp_protocol", "envelopes.py", +) +if os.path.exists(server_protocol_path): + server_pkg = types.ModuleType("comfy_rnp_protocol_e2e_bundle") + server_pkg.__path__ = [os.path.join(_SERVER_DIR, "comfy_rnp_protocol")] + sys.modules["comfy_rnp_protocol_e2e_bundle"] = server_pkg + _load_module( + "comfy_rnp_protocol_e2e_bundle.constants", + os.path.join(_SERVER_DIR, "comfy_rnp_protocol", "constants.py"), + ) + server_envelopes = _load_module( + "comfy_rnp_protocol_e2e_bundle.envelopes", + os.path.join(_SERVER_DIR, "comfy_rnp_protocol", "envelopes.py"), + ) + assert hasattr(server_envelopes, "make_model3d_bundle_envelope"), ( + "Server envelopes module missing make_model3d_bundle_envelope — " + "server-side PR #38 should have added it." + ) + server_env = server_envelopes.make_model3d_bundle_envelope( + files=[ + {"path": "model.obj", "format": "obj", "role": "mesh", + "data": base64.b64encode(obj_bytes).decode("ascii"), + "byte_size": len(obj_bytes)}, + {"path": "model.mtl", "format": "mtl", "role": "material", + "data": base64.b64encode(mtl_bytes).decode("ascii")}, + {"path": "textures/albedo.png", "format": "png", "role": "texture_diffuse", + "data": base64.b64encode(albedo_bytes).decode("ascii")}, + ], + primary_path="model.obj", + format="obj", + ) + assert server_env["type"] == "model_3d", server_env + assert server_env["encoding"] == "bundle_inline", server_env + assert server_env["primary_path"] == "model.obj" + assert len(server_env["files"]) == 3 + # Path sanitization happens server-side. Confirm a malicious path + # would be rejected at envelope build time (defence-in-depth — the + # client decoder doesn't see this). + try: + server_envelopes.make_model3d_bundle_envelope( + files=[{"path": "../etc/passwd", "data": "AA=="}], + primary_path="../etc/passwd", + ) + except ValueError as e: + assert "parent-dir" in str(e) or "absolute" in str(e), ( + f"path sanitization mistake: {e!r}" + ) + else: + raise AssertionError("server should reject '..' path at envelope build time") + # Pipe through the client decoder. + server_bundle = serialization.decode_model3d_envelope(server_env) + assert isinstance(server_bundle, bundle_cls), ( + f"server-built envelope decoded to wrong type: {type(server_bundle)}" + ) + assert server_bundle.get_bytes() == obj_bytes + assert server_bundle.file_bytes("textures/albedo.png") == albedo_bytes + assert server_bundle.file_role("model.obj") == "mesh" + print( + "PASS: end-to-end server make_model3d_bundle_envelope -> client " + "decode_model3d_envelope preserves all files + roles + bytes." + ) +else: + print( + f"SKIP end-to-end: server worktree not found at {_SERVER_DIR!r}" + ) + + +# --------------------------------------------------------------------------- +# 9. CLIENT_CAPABILITIES advertises model_3d:bundle_inline. +# 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.MODEL_3D_BUNDLE_INLINE" in cc_block, ( + "Capability.MODEL_3D_BUNDLE_INLINE not advertised in CLIENT_CAPABILITIES" +) +print("PASS: CLIENT_CAPABILITIES advertises Capability.MODEL_3D_BUNDLE_INLINE.") + +print("ALL CHECKS PASSED.") diff --git a/protocol.py b/protocol.py index 1cb3d8c..086ebe7 100644 --- a/protocol.py +++ b/protocol.py @@ -164,6 +164,67 @@ class Capability: # parse — the magic-byte check at offset 0 is the only validity # guard the server applies). MODEL_3D_GLB_INLINE = "model_3d:glb_inline" + # 3D MODEL BUNDLE envelope: a primary mesh file plus N companion + # files (textures / .mtl / .bin / other PBR maps) carried inline as + # a single coherent unit. Wire shape: + # + # { + # "type": "model_3d", + # "encoding": "bundle_inline", + # "format": "obj", + # "primary_path": "model.obj", + # "files": [ + # {"path": "model.obj", "format": "obj", "role": "mesh", + # "data": "", "byte_size": 12345}, + # {"path": "model.mtl", "format": "mtl", "role": "material", + # "data": ""}, + # {"path": "textures/albedo.png", "format": "png", "role": "texture_diffuse", + # "data": ""} + # ] + # } + # + # ``path`` is the authoritative cross-file reference key — a POSIX- + # relative path (no ``..``, no absolute prefix, no Windows drive + # letter, no backslashes, no case-insensitive collisions). The + # client decoder materialises files into a temp directory using + # exactly these paths so a ``.obj`` that references its ``.mtl`` + # via filename (or a ``.gltf`` that references its ``.bin`` / + # textures via relative URI) resolves on disk. + # + # ``role`` is an OPTIONAL producer hint (NOT a closed enum) — + # well-known values include ``mesh``, ``material``, ``buffer``, + # ``texture_diffuse``, ``texture_metallic``, ``texture_normal``, + # ``texture_roughness``, ``texture_ao``, ``texture_emissive``, + # ``texture_height``. Unknown roles round-trip without error; + # consumers may use ``role`` for convenience but ``path`` is the + # primary identifier. + # + # ``primary_path`` MUST appear in ``files`` exactly once and pins + # the entry-point file (the one passed to ComfyUI's + # ``Types.File3D`` constructor by the client decoder). The + # top-level ``format`` (if set) MUST match that primary file's + # ``format``. + # + # The client decoder returns a ``BundledFile3D`` (a thin + # ``File3D``-compatible wrapper) — existing downstream loaders + # that call ``get_source()`` / ``get_data()`` / ``save_to()`` get + # the bundle's primary file with all sidecars materialised on + # disk first so cross-file refs resolve correctly. Override of + # ``save_to()`` copies the *whole bundle* into the destination + # directory, not just the primary file. + # + # Negotiation: the server emits this encoding ONLY when the + # inbound request advertised this token; legacy clients without + # the bundle decoder will surface NEGOTIATION_FAILED at + # descriptor-load time. Whole-descriptor gating — if a node has + # both bundle and non-bundle outputs and the cap is absent, the + # whole descriptor is skipped rather than synthesising a + # partial-socket variant. + # + # A future sibling ``Capability.MODEL_3D_BUNDLE_URI`` would carry + # ``uri`` per file for URL-fetch-out-of-band semantics; not in + # this PR. + MODEL_3D_BUNDLE_INLINE = "model_3d:bundle_inline" HEAVY_TYPES = frozenset({"image", "video", "audio", "mask", "model_3d"}) diff --git a/serialization.py b/serialization.py index c721be7..c1420cb 100644 --- a/serialization.py +++ b/serialization.py @@ -484,27 +484,315 @@ def decode_video_envelope(envelope: dict[str, Any]) -> Any: # 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. - Mirrors :func:`decode_video_envelope`: the only currently supported - encoding is ``glb_inline`` (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. + 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": - raise RnpProtocolError( - f"Unsupported model_3d encoding: {encoding!r}", - code=ErrorCode.INTERNAL, + 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, ) - from comfy_api.latest._util import File3D - file_format = envelope.get("format") or "glb" - return File3D(BytesIO(decode_envelope_data(envelope)), str(file_format)) + raise RnpProtocolError( + f"Unsupported model_3d encoding: {encoding!r}", + code=ErrorCode.INTERNAL, + ) # --------------------------------------------------------------------------- @@ -567,6 +855,7 @@ __all__ = [ "decode_audio_envelope", "decode_video_envelope", "decode_model3d_envelope", + "_bundled_file3d_class", "decode_envelope", "is_envelope", "is_image_tensor",