From 5bb76c13b114b467bd106f5c9c8e098cc60451c3 Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Tue, 19 May 2026 10:22:23 -0400 Subject: [PATCH] CORE-217: Add IMAGE/MASK outputs to Preview3D and SaveGLB --- app/preview3d_bridge.py | 116 ++++++++++++++++++++++++++++++++++ comfy_extras/nodes_load_3d.py | 26 +++++++- comfy_extras/nodes_save_3d.py | 30 ++++++++- server.py | 3 + 4 files changed, 169 insertions(+), 6 deletions(-) create mode 100644 app/preview3d_bridge.py diff --git a/app/preview3d_bridge.py b/app/preview3d_bridge.py new file mode 100644 index 000000000..d357f51f6 --- /dev/null +++ b/app/preview3d_bridge.py @@ -0,0 +1,116 @@ +"""Round-trip render bridge for Preview3D / SaveGLB IMAGE+MASK outputs. + +Flow: a node's execute() calls bridge.request_render(), which sends a +websocket "preview3d.render_request" event to the client and awaits a +Future. The client renders the file with a fixed camera at fixed +resolution, uploads the PNGs to /temp/, and POSTs {render_id, image, +mask} to /3d/render_response, which resolves the Future. The node then +loads image/mask via the standard LoadImage path. +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import uuid +from typing import Any + +from aiohttp import web + + +_RENDER_TIMEOUT_SEC = 60 + + +class Preview3DBridge: + """Server-side half of the Preview3D / SaveGLB render round-trip. + + Owns the pending Future registry and the POST /3d/render_response + route. Instantiated once on PromptServer and surfaced as + PromptServer.instance.preview3d_bridge. + """ + + def __init__(self) -> None: + self._pending: dict[str, concurrent.futures.Future] = {} + + def add_routes(self, routes: web.RouteTableDef) -> None: + @routes.post("/3d/render_response") + async def render_response(request: web.Request) -> web.Response: + try: + data = await request.json() + except Exception: + return web.Response(status=400, text="invalid json") + render_id = data.get("render_id") + if not render_id or render_id not in self._pending: + return web.Response(status=404, text="unknown render_id") + future = self._pending.pop(render_id) + if future.done(): + return web.Response(status=409, text="already resolved") + if "error" in data: + future.set_exception(RuntimeError(str(data["error"]))) + else: + future.set_result({ + "image": data.get("image"), + "mask": data.get("mask"), + }) + return web.json_response({"ok": True}) + + async def request_render( + self, + node_id: str, + file_path: str, + file_type: str = "output", + camera_info: dict | None = None, + timeout: float = _RENDER_TIMEOUT_SEC, + ) -> dict[str, Any]: + """Send a render request and await the client's response. + + node_id should be the executor-provided cls.hidden.unique_id of + the requesting Preview3D / SaveGLB node — the client uses it to + find the matching viewer instance. + """ + from server import PromptServer + server = PromptServer.instance + if server is None: + raise RuntimeError("PromptServer is not initialized") + + render_id = uuid.uuid4().hex + future: concurrent.futures.Future = concurrent.futures.Future() + self._pending[render_id] = future + + payload = { + "render_id": render_id, + "node_id": node_id, + "file_path": file_path, + "type": file_type, + "camera_info": camera_info, + } + server.send_sync("preview3d.render_request", payload, server.client_id) + + try: + return await asyncio.wait_for( + asyncio.wrap_future(future), timeout=timeout + ) + except asyncio.TimeoutError: + self._pending.pop(render_id, None) + raise RuntimeError( + f"Preview3D render bridge: client did not respond in {timeout}s " + f"(node_id={node_id}, file={file_path}). Is the workflow open " + "in a browser tab?" + ) + except Exception: + self._pending.pop(render_id, None) + raise + + +def load_image_and_mask(image_ref: str, mask_ref: str): + import nodes + loader = nodes.LoadImage() + image_tensor, _ = loader.load_image(image=image_ref) + _, mask_tensor = loader.load_image(image=mask_ref) + return image_tensor, mask_tensor + + +__all__ = [ + "Preview3DBridge", + "load_image_and_mask", +] diff --git a/comfy_extras/nodes_load_3d.py b/comfy_extras/nodes_load_3d.py index 9112bdd0a..b01baff12 100644 --- a/comfy_extras/nodes_load_3d.py +++ b/comfy_extras/nodes_load_3d.py @@ -6,6 +6,9 @@ import uuid from typing_extensions import override from comfy_api.latest import IO, UI, ComfyExtension, InputImpl, Types +from app.preview3d_bridge import load_image_and_mask +from server import PromptServer + from pathlib import Path @@ -101,11 +104,15 @@ class Preview3D(IO.ComfyNode): IO.Load3DCamera.Input("camera_info", optional=True, advanced=True), IO.Image.Input("bg_image", optional=True, advanced=True), ], - outputs=[], + outputs=[ + IO.Image.Output(display_name="image"), + IO.Mask.Output(display_name="mask"), + ], + hidden=[IO.Hidden.unique_id], ) @classmethod - def execute(cls, model_file: str | Types.File3D, **kwargs) -> IO.NodeOutput: + async def execute(cls, model_file: str | Types.File3D, **kwargs) -> IO.NodeOutput: if isinstance(model_file, Types.File3D): filename = f"preview3d_{uuid.uuid4().hex}.{model_file.format}" model_file.save_to(os.path.join(folder_paths.get_output_directory(), filename)) @@ -113,7 +120,20 @@ class Preview3D(IO.ComfyNode): filename = model_file camera_info = kwargs.get("camera_info", None) bg_image = kwargs.get("bg_image", None) - return IO.NodeOutput(ui=UI.PreviewUI3D(filename, camera_info, bg_image=bg_image)) + + rendered = await PromptServer.instance.preview3d_bridge.request_render( + node_id=cls.hidden.unique_id, + file_path=filename, + file_type="output", + camera_info=camera_info, + ) + image_tensor, mask_tensor = load_image_and_mask(rendered["image"], rendered["mask"]) + + return IO.NodeOutput( + image_tensor, + mask_tensor, + ui=UI.PreviewUI3D(filename, camera_info, bg_image=bg_image), + ) process = execute # TODO: remove diff --git a/comfy_extras/nodes_save_3d.py b/comfy_extras/nodes_save_3d.py index c03524246..90b4ccba8 100644 --- a/comfy_extras/nodes_save_3d.py +++ b/comfy_extras/nodes_save_3d.py @@ -14,6 +14,8 @@ from typing_extensions import override import folder_paths from comfy.cli_args import args from comfy_api.latest import ComfyExtension, IO, Types +from app.preview3d_bridge import load_image_and_mask +from server import PromptServer def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None): @@ -329,12 +331,17 @@ class SaveGLB(IO.ComfyNode): tooltip="Mesh or 3D file to save", ), IO.String.Input("filename_prefix", default="3d/ComfyUI"), + IO.Load3DCamera.Input("camera_info", optional=True, advanced=True), ], - hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo] + outputs=[ + IO.Image.Output(display_name="image"), + IO.Mask.Output(display_name="mask"), + ], + hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo, IO.Hidden.unique_id] ) @classmethod - def execute(cls, mesh: Types.MESH | Types.File3D, filename_prefix: str) -> IO.NodeOutput: + async def execute(cls, mesh: Types.MESH | Types.File3D, filename_prefix: str, camera_info=None) -> IO.NodeOutput: full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, folder_paths.get_output_directory()) results = [] @@ -383,7 +390,24 @@ class SaveGLB(IO.ComfyNode): "type": "output" }) counter += 1 - return IO.NodeOutput(ui={"3d": results}) + + if not results: + raise ValueError("SaveGLB: no mesh was written; cannot render image output") + + first = results[0] + rel_path = (first["subfolder"] + "/" + first["filename"]) if first["subfolder"] else first["filename"] + rendered = await PromptServer.instance.preview3d_bridge.request_render( + node_id=cls.hidden.unique_id, + file_path=rel_path, + file_type="output", + camera_info=camera_info, + ) + image_tensor, mask_tensor = load_image_and_mask(rendered["image"], rendered["mask"]) + + ui_out: dict = {"3d": results} + if camera_info is not None: + ui_out["camera_info"] = [camera_info] + return IO.NodeOutput(image_tensor, mask_tensor, ui=ui_out) class Save3DExtension(ComfyExtension): diff --git a/server.py b/server.py index 44470b904..57626b0c4 100644 --- a/server.py +++ b/server.py @@ -44,6 +44,7 @@ from app.model_manager import ModelFileManager from app.custom_node_manager import CustomNodeManager from app.subgraph_manager import SubgraphManager from app.node_replace_manager import NodeReplaceManager +from app.preview3d_bridge import Preview3DBridge from typing import Optional, Union from api_server.routes.internal.internal_routes import InternalRoutes from protocol import BinaryEventTypes @@ -209,6 +210,7 @@ class PromptServer(): self.custom_node_manager = CustomNodeManager() self.subgraph_manager = SubgraphManager() self.node_replace_manager = NodeReplaceManager() + self.preview3d_bridge = Preview3DBridge() self.internal_routes = InternalRoutes(self) self.supports = ["custom_nodes_from_web"] self.prompt_queue = execution.PromptQueue(self) @@ -1050,6 +1052,7 @@ class PromptServer(): self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items()) self.subgraph_manager.add_routes(self.routes, nodes.LOADED_MODULE_DIRS.items()) self.node_replace_manager.add_routes(self.routes) + self.preview3d_bridge.add_routes(self.routes) self.app.add_subapp('/internal', self.internal_routes.get_app()) # Prefix every route with /api for easier matching for delegation.