Support dynamic Gaussian updates

src/viser/_scene_handles.py

@@ -26,7 +26,7 @@
 from typing_extensions import Self, deprecated, override
 
 from . import _messages
-from ._assignable_props_api import AssignablePropsBase
+from ._assignable_props_api import AssignablePropsBase, colors_to_uint8
 from .infra._infra import (
     WebsockClientConnection,
     WebsockServer,
@@ -1235,6 +1235,70 @@ def _ensure_buffer_size(self, num_gaussians: int) -> None:
 
         self.buffer = new_buffer
 
+    @staticmethod
+    def _pack_centers(buffer: np.ndarray, centers: np.ndarray) -> None:
+        buffer[:, 0:3] = np.ascontiguousarray(centers, dtype=np.float32).view(np.uint32)
+
+    @staticmethod
+    def _pack_covariances(buffer: np.ndarray, covariances: np.ndarray) -> None:
+        # Extract upper-triangular terms: indices [0,1,2,4,5,8] from flattened 3x3.
+        cov_triu = covariances.reshape((-1, 9))[:, np.array([0, 1, 2, 4, 5, 8])]
+        cov_triu_f16 = np.ascontiguousarray(cov_triu, dtype=np.float16)
+        buffer[:, 4:7] = cov_triu_f16.view(np.uint32)
+
+    @staticmethod
+    def _pack_rgba(
+        buffer: np.ndarray,
+        rgbs: np.ndarray | None = None,
+        opacities: np.ndarray | None = None,
+    ) -> None:
+        rgba = buffer[:, 7:8].view(np.uint8).reshape(-1, 4)
+        if rgbs is not None:
+            rgba[:, :3] = colors_to_uint8(rgbs)
+        if opacities is not None:
+            rgba[:, 3:4] = colors_to_uint8(opacities)
+        buffer[:, 7:8] = rgba.view(np.uint32)
+
+    def set_gaussians(
+        self,
+        centers: np.ndarray,
+        covariances: np.ndarray,
+        rgbs: np.ndarray,
+        opacities: np.ndarray,
+    ) -> None:
+        """Atomically update all Gaussian attributes, including count changes.
+
+        This is the preferred fast path when per-frame updates may change the
+        number of Gaussians: all attributes land in a single buffer update,
+        preventing transient mixed-state frames from sequential property
+        assignments (centers/covariances/rgbs/opacities one-by-one). A call
+        that leaves the buffer numerically unchanged sends no message.
+        """
+        assert centers.ndim == 2 and centers.shape[1] == 3, (
+            f"centers must have shape (N, 3), got {centers.shape}"
+        )
+        num_gaussians = centers.shape[0]
+        assert covariances.ndim == 3 and covariances.shape == (num_gaussians, 3, 3), (
+            f"covariances must have shape ({num_gaussians}, 3, 3), got {covariances.shape}"
+        )
+        assert rgbs.ndim == 2 and rgbs.shape == (num_gaussians, 3), (
+            f"rgbs must have shape ({num_gaussians}, 3), got {rgbs.shape}"
+        )
+        assert opacities.ndim == 2 and opacities.shape == (num_gaussians, 1), (
+            f"opacities must have shape ({num_gaussians}, 1), got {opacities.shape}"
+        )
+
+        # Assemble the full buffer locally, then store it with a single
+        # property assignment: props_setattr rejects writes to removed handles,
+        # resizes or copies in place as needed, and queues exactly one private
+        # snapshot for the wire. Routing a resize through _ensure_buffer_size
+        # here would queue an extra all-default buffer message first.
+        buffer = np.zeros((num_gaussians, 8), dtype=np.uint32)
+        self._pack_centers(buffer, centers)
+        self._pack_covariances(buffer, covariances)
+        self._pack_rgba(buffer, rgbs=rgbs, opacities=opacities)
+        self.buffer = buffer
+
     @property
     def centers(self) -> npt.NDArray[np.float32]:
         """Centers of the Gaussians. Shape: (N, 3). Synchronized automatically when assigned."""
@@ -1246,7 +1310,7 @@ def centers(self, centers: np.ndarray) -> None:
             f"centers must have shape (N, 3), got {centers.shape}"
         )
         self._ensure_buffer_size(centers.shape[0])
-        self.buffer[:, 0:3] = centers.astype(np.float32).view(np.uint32)
+        self._pack_centers(self.buffer, centers)
         # Queue a private snapshot: the stored buffer is mutated in place by
         # later sub-property assignments, possibly while the event loop is still
         # serializing this message. Matches the guard in props_setattr.
@@ -1260,15 +1324,11 @@ def rgbs(self) -> npt.NDArray[np.uint8]:
 
     @rgbs.setter
     def rgbs(self, rgbs: np.ndarray) -> None:
-        from ._assignable_props_api import colors_to_uint8
-
         assert rgbs.ndim == 2 and rgbs.shape[1] == 3, (
             f"rgbs must have shape (N, 3), got {rgbs.shape}"
         )
         self._ensure_buffer_size(rgbs.shape[0])
-        rgba = self.buffer[:, 7:8].view(np.uint8).reshape(-1, 4)
-        rgba[:, :3] = colors_to_uint8(rgbs)
-        self.buffer[:, 7:8] = rgba.view(np.uint32)
+        self._pack_rgba(self.buffer, rgbs=rgbs)
         self._queue_update("buffer", self.buffer.copy())
 
     @property
@@ -1280,15 +1340,11 @@ def opacities(self) -> npt.NDArray[np.uint8]:
 
     @opacities.setter
     def opacities(self, opacities: np.ndarray) -> None:
-        from ._assignable_props_api import colors_to_uint8
-
         assert opacities.ndim == 2 and opacities.shape[1] == 1, (
             f"opacities must have shape (N, 1), got {opacities.shape}"
         )
         self._ensure_buffer_size(opacities.shape[0])
-        rgba = self.buffer[:, 7:8].view(np.uint8).reshape(-1, 4)
-        rgba[:, 3:4] = colors_to_uint8(opacities)
-        self.buffer[:, 7:8] = rgba.view(np.uint32)
+        self._pack_rgba(self.buffer, opacities=opacities)
         self._queue_update("buffer", self.buffer.copy())
 
     @property
@@ -1317,10 +1373,7 @@ def covariances(self, covariances: np.ndarray) -> None:
             f"covariances must have shape (N, 3, 3), got {covariances.shape}"
         )
         self._ensure_buffer_size(covariances.shape[0])
-        # Extract upper-triangular terms: indices [0,1,2,4,5,8] from flattened 3x3.
-        cov_triu = covariances.reshape((-1, 9))[:, np.array([0, 1, 2, 4, 5, 8])]
-        cov_triu_f16 = cov_triu.astype(np.float16)
-        self.buffer[:, 4:7] = np.ascontiguousarray(cov_triu_f16).view(np.uint32)
+        self._pack_covariances(self.buffer, covariances)
         self._queue_update("buffer", self.buffer.copy())
 
 

tests/test_handle_lifecycle_bugs.py

@@ -448,6 +448,107 @@ def latest_buffer_update() -> np.ndarray:
         assert np.array_equal(queued, before)
 
 
+def _add_test_splats(server: viser.ViserServer, n: int):
+    return server.scene.add_gaussian_splats(
+        "/splats",
+        centers=np.zeros((n, 3), np.float32),
+        covariances=np.tile(np.eye(3, dtype=np.float32) * 0.01, (n, 1, 1)),
+        rgbs=np.zeros((n, 3), np.uint8),
+        opacities=np.ones((n, 1), np.float32),
+    )
+
+
+def test_set_gaussians_round_trip() -> None:
+    """set_gaussians() must store values readable back through the per-property
+    getters, including when the Gaussian count changes."""
+    with _server() as server:
+        s = _add_test_splats(server, n=2)
+
+        m = 4  # Different count from the initial 2.
+        centers = np.arange(m * 3, dtype=np.float32).reshape(m, 3)
+        covariances = np.tile(np.eye(3, dtype=np.float32) * 0.25, (m, 1, 1))
+        rgbs = np.arange(m * 3, dtype=np.uint8).reshape(m, 3)
+        opacities = np.linspace(0.0, 1.0, m, dtype=np.float32).reshape(m, 1)
+        s.set_gaussians(centers, covariances, rgbs, opacities)
+
+        assert np.array_equal(s.centers, centers)
+        assert np.array_equal(s.rgbs, rgbs)
+        # Opacities are quantized to uint8 on the way in.
+        assert np.array_equal(
+            s.opacities, np.clip(opacities * 255.0, 0, 255).astype(np.uint8)
+        )
+        # Covariances round-trip through float16.
+        np.testing.assert_allclose(s.covariances, covariances, rtol=1e-3)
+
+        # Zero Gaussians is a valid state.
+        s.set_gaussians(
+            np.zeros((0, 3), np.float32),
+            np.zeros((0, 3, 3), np.float32),
+            np.zeros((0, 3), np.uint8),
+            np.zeros((0, 1), np.float32),
+        )
+        assert s.centers.shape == (0, 3)
+
+
+def test_set_gaussians_count_change_queues_single_update() -> None:
+    """A count-changing set_gaussians() must push exactly one buffer update: an
+    intermediate resize message would flash an all-default buffer on clients
+    whenever the ~60 Hz flush lands between the two pushes."""
+    with _server() as server:
+        s = _add_test_splats(server, n=2)
+
+        interface = s._impl.api._websock_interface
+        original_queue_message = interface.queue_message
+        buffer_updates: list[np.ndarray] = []
+
+        def spy(message) -> None:
+            if (
+                type(message).__name__ == "SceneNodeUpdateMessage"
+                and "buffer" in message.updates
+            ):
+                buffer_updates.append(message.updates["buffer"])
+            original_queue_message(message)
+
+        interface.queue_message = spy  # type: ignore[method-assign]
+        try:
+            m = 5
+            s.set_gaussians(
+                np.ones((m, 3), np.float32),
+                np.tile(np.eye(3, dtype=np.float32) * 0.01, (m, 1, 1)),
+                np.zeros((m, 3), np.uint8),
+                np.ones((m, 1), np.float32),
+            )
+        finally:
+            interface.queue_message = original_queue_message  # type: ignore[method-assign]
+
+        assert len(buffer_updates) == 1
+        # The queued snapshot must alias neither the live server-owned buffer
+        # (mutated in place by later sub-property writes) nor anything the
+        # caller can touch.
+        assert buffer_updates[0] is not s._impl.props.buffer
+        queued_before = buffer_updates[0].copy()
+        s.centers = np.full((m, 3), 7.0, np.float32)
+        assert np.array_equal(buffer_updates[0], queued_before)
+
+
+def test_set_gaussians_rejects_removed_handle() -> None:
+    """set_gaussians() on a removed handle must raise like property assignment
+    does, instead of queuing updates that resurrect the node on clients."""
+    with _server() as server:
+        n = 3
+        s = _add_test_splats(server, n=n)
+        s.remove()
+
+        for count in (n, n + 1):  # Same-count and resizing calls alike.
+            with pytest.raises(RuntimeError):
+                s.set_gaussians(
+                    np.zeros((count, 3), np.float32),
+                    np.tile(np.eye(3, dtype=np.float32) * 0.01, (count, 1, 1)),
+                    np.zeros((count, 3), np.uint8),
+                    np.ones((count, 1), np.float32),
+                )
+
+
 def test_upload_part_after_button_removal_still_acks_and_completes() -> None:
     """A FileTransferPart arriving after button.remove() mid-upload must not
     assert (the handle is legitimately gone from the registry): parts keep