CORE-217: Add IMAGE/MASK outputs to Preview3D and SaveGLB

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",
+]

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
 

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):

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.