Support multiclient in viewer_beta with separate render threads for each client (#2519)

* Initial multiclient, jittery screen bc fighting between clients

* Fix multiclient jitter with separate render threads for each client

* Handle client disconnect

* Fix viewer.init_scene and viewer_elements.get_camera for multiclient

---------

Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>
Co-authored-by: Brent Yi <yibrenth@gmail.com>

nerfstudio/viewer_beta/render_state_machine.py

@@ -21,6 +21,7 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, get_args
 
 import torch
+from viser import ClientHandle
 from nerfstudio.model_components.renderers import background_color_override_context
 from nerfstudio.utils import colormaps, writer
 from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
@@ -52,7 +53,7 @@ class RenderStateMachine(threading.Thread):
         viewer: the viewer state
     """
 
-    def __init__(self, viewer: Viewer, viser_scale_ratio: float):
+    def __init__(self, viewer: Viewer, viser_scale_ratio: float, client: ClientHandle):
         threading.Thread.__init__(self)
         self.transitions: Dict[RenderStates, Dict[RenderActions, RenderStates]] = {
             s: {} for s in get_args(RenderStates)
@@ -77,6 +78,8 @@ class RenderStateMachine(threading.Thread):
         self.daemon = True
         self.output_keys = {}
         self.viser_scale_ratio = viser_scale_ratio
+        self.client = client
+        self.running = True
 
     def action(self, action: RenderAction):
         """Takes an action and updates the state machine
@@ -107,7 +110,7 @@ class RenderStateMachine(threading.Thread):
         self.render_trigger.set()
 
     def _render_img(self, camera_state: CameraState):
-        """Takes the current camera, generates rays, and renders the iamge
+        """Takes the current camera, generates rays, and renders the image
 
         Args:
             camera_state: the current camera state
@@ -166,7 +169,7 @@ class RenderStateMachine(threading.Thread):
 
     def run(self):
         """Main loop for the render thread"""
-        while True:
+        while self.running:
             if not self.render_trigger.wait(0.2):
                 # if we haven't received a trigger in a while, send a static action
                 if self.viewer.camera_state is not None:
@@ -237,7 +240,8 @@ class RenderStateMachine(threading.Thread):
         depth = (
             outputs["gl_z_buf_depth"].cpu().numpy() * self.viser_scale_ratio if "gl_z_buf_depth" in outputs else None
         )
-        self.viewer.viser_server.set_background_image(
+
+        self.client.set_background_image(
             selected_output.cpu().numpy(),
             format=self.viewer.config.image_format,
             jpeg_quality=self.viewer.config.jpeg_quality,

nerfstudio/viewer_beta/viewer.py

@@ -103,8 +103,8 @@ class Viewer:
         self.step = 0
         self.train_btn_state: Literal["training", "paused", "completed"] = "training"
         self._prev_train_state: Literal["training", "paused", "completed"] = "training"
+        self.last_move_time = 0
 
-        self.client: Optional[viser.ClientHandle] = None
         self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port, share=share)
         buttons = (
             viser.theme.TitlebarButton(
@@ -137,7 +137,8 @@ class Viewer:
             brand_color=(255, 211, 105),
         )
 
-        self.render_statemachine = RenderStateMachine(self, VISER_NERFSTUDIO_SCALE_RATIO)
+        self.render_statemachines: Dict[int, RenderStateMachine] = {}
+        self.viser_server.on_client_disconnect(self.handle_disconnect)
         self.viser_server.on_client_connect(self.handle_new_client)
 
         tabs = self.viser_server.add_gui_tab_group()
@@ -202,24 +203,30 @@ class Viewer:
             ]
         for c in self.viewer_controls:
             c._setup(self)
-        self.render_statemachine.start()
+
+    def get_camera_state(self, client: viser.ClientHandle) -> CameraState:
+        R = vtf.SO3(wxyz=client.camera.wxyz)
+        R = R @ vtf.SO3.from_x_radians(np.pi)
+        R = torch.tensor(R.as_matrix())
+        pos = torch.tensor(client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
+        c2w = torch.concatenate([R, pos[:, None]], dim=1)
+        camera_state = CameraState(fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w)
+        return camera_state
+
+    def handle_disconnect(self, client: viser.ClientHandle) -> None:
+        self.render_statemachines[client.client_id].running = False
+        self.render_statemachines.pop(client.client_id)
 
     def handle_new_client(self, client: viser.ClientHandle) -> None:
-        self.client = client
-        self.last_move_time = 0
+        self.render_statemachines[client.client_id] = RenderStateMachine(self, VISER_NERFSTUDIO_SCALE_RATIO, client)
+        self.render_statemachines[client.client_id].start()
 
         @client.camera.on_update
         def _(cam: viser.CameraHandle) -> None:
-            assert self.client is not None
-            with client.atomic():
-                self.last_move_time = time.time()
-                R = vtf.SO3(wxyz=self.client.camera.wxyz)
-                R = R @ vtf.SO3.from_x_radians(np.pi)
-                R = torch.tensor(R.as_matrix())
-                pos = torch.tensor(self.client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
-                c2w = torch.concatenate([R, pos[:, None]], dim=1)
-                self.camera_state = CameraState(fov=self.client.camera.fov, aspect=self.client.camera.aspect, c2w=c2w)
-                self.render_statemachine.action(RenderAction("move", self.camera_state))
+            self.last_move_time = time.time()
+            with self.viser_server.atomic():
+                camera_state = self.get_camera_state(client)
+                self.render_statemachines[client.client_id].action(RenderAction("move", camera_state))
 
     def set_camera_visibility(self, visible: bool) -> None:
         """Toggle the visibility of the training cameras."""
@@ -250,8 +257,11 @@ class Viewer:
 
     def _interrupt_render(self, _) -> None:
         """Interrupt current render."""
-        if self.camera_state is not None:
-            self.render_statemachine.action(RenderAction("rerender", self.camera_state))
+        clients = self.viser_server.get_clients()
+        for id in clients:
+            camera_state = self.get_camera_state(clients[id])
+            if camera_state is not None:
+                self.render_statemachines[id].action(RenderAction("rerender", camera_state))
 
     def _toggle_training_state(self, _) -> None:
         """Toggle the trainer's training state."""
@@ -263,8 +273,11 @@ class Viewer:
 
     def _crop_params_update(self, _) -> None:
         """Update crop parameters"""
-        if self.camera_state is not None:
-            self.render_statemachine.action(RenderAction("move", self.camera_state))
+        clients = self.viser_server.get_clients()
+        for id in clients:
+            camera_state = self.get_camera_state(clients[id])
+            if camera_state is not None:
+                self.render_statemachines[id].action(RenderAction("move", camera_state))
 
     def _output_type_change(self, _):
         self.output_type_changed = True
@@ -327,10 +340,9 @@ class Viewer:
 
             @camera_handle.on_click
             def _(event: viser.SceneNodePointerEvent[viser.CameraFrustumHandle]) -> None:
-                assert self.client is not None
-                with self.client.atomic():
-                    self.client.camera.position = event.target.position
-                    self.client.camera.wxyz = event.target.wxyz
+                with event.client.atomic():
+                    event.client.camera.position = event.target.position
+                    event.client.camera.wxyz = event.target.wxyz
 
             self.camera_handles[idx] = camera_handle
             self.original_c2w[idx] = c2w
@@ -347,7 +359,7 @@ class Viewer:
         """
         self.step = step
 
-        if self.camera_state is None:
+        if len(self.render_statemachines) == 0:
             return
         # this stops training while moving to make the response smoother
         while time.time() - self.last_move_time < 0.1:
@@ -371,7 +383,11 @@ class Viewer:
                 render_freq = 30
             if step > self.last_step + render_freq:
                 self.last_step = step
-                self.render_statemachine.action(RenderAction("step", self.camera_state))
+                clients = self.viser_server.get_clients()
+                for id in clients:
+                    camera_state = self.get_camera_state(clients[id])
+                    if camera_state is not None:
+                        self.render_statemachines[id].action(RenderAction("step", camera_state))
                 self.update_camera_poses()
                 self.control_panel.update_step(step)
 

nerfstudio/viewer_beta/viewer_elements.py

@@ -117,7 +117,7 @@ class ViewerControl:
         """
         raise NotImplementedError()
 
-    def get_camera(self, img_height: int, img_width: int) -> Optional[Cameras]:
+    def get_camera(self, img_height: int, img_width: int, client_id: Optional[int] = None) -> Optional[Cameras]:
         """
         Returns the Cameras object representing the current camera for the viewer, or None if the viewer
         is not connected yet
@@ -126,16 +126,21 @@ class ViewerControl:
             img_height: The height of the image to get camera intrinsics for
             img_width: The width of the image to get camera intrinsics for
         """
-        assert self.viewer.client is not None
+        clients = self.viser_server.get_clients()
+        if len(clients) == 0:
+            return None
+        if not client_id:
+            client_id = list(clients.keys())[0]
 
         from nerfstudio.viewer_beta.viewer import VISER_NERFSTUDIO_SCALE_RATIO
 
-        R = vtf.SO3(wxyz=self.viewer.client.camera.wxyz)
+        client = clients[client_id]
+        R = vtf.SO3(wxyz=client.camera.wxyz)
         R = R @ vtf.SO3.from_x_radians(np.pi)
         R = torch.tensor(R.as_matrix())
-        pos = torch.tensor(self.viewer.client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
+        pos = torch.tensor(client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
         c2w = torch.concatenate([R, pos[:, None]], dim=1)
-        camera_state = CameraState(fov=self.viewer.client.camera.fov, aspect=self.viewer.client.camera.aspect, c2w=c2w)
+        camera_state = CameraState(fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w)
         return get_camera(camera_state, img_height, img_width)
 
     def register_click_cb(self, cb: Callable):