@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-"""Demo: in-memory 3D image viewed as 2D slice and 3D volume.
+"""Demo: in-memory 3D image with toggleable 2D/3D view in a single scene.
 
 Run with:
     uv run python scripts/v2/image_memory_demo.py
@@ -14,8 +14,10 @@
 from PySide6.QtCore import QTimer
 from PySide6.QtWidgets import (
     QApplication,
+    QCheckBox,
     QHBoxLayout,
     QLabel,
+    QPushButton,
     QSpinBox,
     QVBoxLayout,
     QWidget,
@@ -32,120 +34,123 @@ def main() -> None:
     app = QApplication(sys.argv)
 
     # ── 1. Create data ──────────────────────────────────────────────────
-    volume = binary_blobs(length=128, volume_fraction=0.1, n_dim=3).astype(np.float32)
+    volume = binary_blobs(length=128, volume_fraction=0.03, n_dim=3).astype(np.float32)
     store = ImageMemoryStore(data=volume, name="demo_volume")
 
     # ── 2. Set up the controller ────────────────────────────────────────
     controller = CellierController()
 
-    # Coordinate system: 3 spatial axes
     cs = CoordinateSystem(name="world", axis_labels=("z", "y", "x"))
 
-    # 2D scene: displayed_axes=(1,2), slice_indices={0:0}
-    scene_2d = controller.add_scene(dim="2d", coordinate_system=cs, name="slice")
-    # 3D scene: displayed_axes=(0,1,2), slice_indices={}
-    scene_3d = controller.add_scene(dim="3d", coordinate_system=cs, name="volume")
-
-    appearance = ImageMemoryAppearance(
-        color_map="viridis",
-        clim=(0.0, 1.0),
+    # Single scene that supports both 2D and 3D rendering.
+    # Start in 3D; displayed_axes=(0,1,2).
+    scene = controller.add_scene(
+        dim="3d",
+        coordinate_system=cs,
+        name="volume",
+        render_modes={"2d", "3d"},
     )
 
-    # Add the same store to both scenes (store is registered only once)
-    visual_2d = controller.add_image(
-        data=store,
-        scene_id=scene_2d.id,
-        appearance=appearance,
-        name="slice_view",
-    )
-    visual_3d = controller.add_image(
+    appearance = ImageMemoryAppearance(color_map="viridis", clim=(0.0, 1.0))
+    visual = controller.add_image(
         data=store,
-        scene_id=scene_3d.id,
+        scene_id=scene.id,
         appearance=appearance,
         name="volume_view",
     )
+    visual.aabb.color = "#ff00ff"
 
-    # ── 3. Create canvases and z-slice control ─────────────────────────
+    # ── 3. Build UI ─────────────────────────────────────────────────────
     root = QWidget()
     outer_layout = QVBoxLayout(root)
     controller.set_widget_parent(root)
 
-    # Canvas row
-    canvas_row = QWidget()
-    canvas_layout = QHBoxLayout(canvas_row)
-    canvas_layout.setContentsMargins(0, 0, 0, 0)
+    canvas_widget = controller.add_canvas(scene_id=scene.id)
+    outer_layout.addWidget(canvas_widget)
 
-    widget_2d = controller.add_canvas(scene_id=scene_2d.id)
-    widget_3d = controller.add_canvas(scene_id=scene_3d.id)
-
-    canvas_layout.addWidget(widget_2d)
-    canvas_layout.addWidget(widget_3d)
-    outer_layout.addWidget(canvas_row)
-
-    # Z-slice spinbox
-    z_depth = store.shape[0]  # axis 0 = z
+    # Controls row
     controls = QWidget()
     ctrl_layout = QHBoxLayout(controls)
     ctrl_layout.setContentsMargins(0, 0, 0, 0)
-    ctrl_layout.addWidget(QLabel("Z slice:"))
 
+    z_depth = store.shape[0]
+    ctrl_layout.addWidget(QLabel("Z slice:"))
     z_spin = QSpinBox()
     z_spin.setRange(0, z_depth - 1)
-    z_spin.setValue(0)
+    z_spin.setValue(z_depth // 2)
     ctrl_layout.addWidget(z_spin)
     ctrl_layout.addStretch()
+
+    aabb_check = QCheckBox("Show bounding box")
+    aabb_check.setChecked(False)
+    ctrl_layout.addWidget(aabb_check)
+
+    toggle_btn = QPushButton("Switch to 2D")
+    ctrl_layout.addWidget(toggle_btn)
     outer_layout.addWidget(controls)
 
+    root.resize(900, 700)
+    root.show()
+
+    # ── 4. Toggle logic ─────────────────────────────────────────────────
+    _current_dim = ["3d"]
+
+    def _toggle() -> None:
+        if _current_dim[0] == "3d":
+            _current_dim[0] = "2d"
+            toggle_btn.setText("Switch to 3D")
+            scene.dims.selection.slice_indices = {0: z_spin.value()}
+            scene.dims.selection.displayed_axes = (1, 2)
+        else:
+            _current_dim[0] = "3d"
+            toggle_btn.setText("Switch to 2D")
+            scene.dims.selection.displayed_axes = (0, 1, 2)
+
+        controller.reslice_scene(scene.id)
+
+    toggle_btn.clicked.connect(_toggle)
+
     def _on_z_changed(z: int) -> None:
-        scene_2d.dims.selection.slice_indices = {0: z}
-        controller.reslice_scene(scene_2d.id)
+        scene.dims.selection.slice_indices = {0: z}
+        if _current_dim[0] == "2d":
+            controller.reslice_scene(scene.id)
 
     z_spin.valueChanged.connect(_on_z_changed)
 
-    root.resize(1200, 700)
-    root.show()
+    def _on_aabb_toggled(checked: bool) -> None:
+        visual.aabb.enabled = checked
 
-    # ── 4. Initial reslice + camera fit ────────────────────────────────
-    # The async slicer needs a running event loop, so reslice is deferred
-    # with QTimer.singleShot(0, ...) until QtAsyncio.run() has started.
-    #
-    # look_at_visual computes the bounding box from the current node
-    # geometry, so it must run *after* data has been committed to the GPU
-    # (otherwise the placeholder texture gives a near-zero bounding box).
-    # We wrap the on_data_ready callbacks to trigger camera fit on the
-    # first delivery — event-driven, not time-based.
-
-    visuals_loaded: set[str] = set()
-
-    def _wrap_callback(original, visual_id, label, view_dir, up):
-        def _wrapped(batch):
-            original(batch)
-            if label not in visuals_loaded:
-                visuals_loaded.add(label)
-                controller.look_at_visual(visual_id, view_direction=view_dir, up=up)
-
-        return _wrapped
-
-    # Patch the render-layer visuals to fire camera fit on first data.
-    scene_mgr_2d = controller._render_manager._scenes[scene_2d.id]
-    gfx_vis_2d = scene_mgr_2d.get_visual(visual_2d.id)
-    gfx_vis_2d.on_data_ready_2d = _wrap_callback(
-        gfx_vis_2d.on_data_ready_2d,
-        visual_2d.id,
-        "2d",
-        view_dir=(0, 0, -1),
-        up=(0, 1, 0),
-    )
+    aabb_check.toggled.connect(_on_aabb_toggled)
 
-    scene_mgr_3d = controller._render_manager._scenes[scene_3d.id]
-    gfx_vis_3d = scene_mgr_3d.get_visual(visual_3d.id)
-    gfx_vis_3d.on_data_ready = _wrap_callback(
-        gfx_vis_3d.on_data_ready,
-        visual_3d.id,
-        "3d",
-        view_dir=(-1, -1, -1),
-        up=(0, 0, 1),
-    )
+    # ── 5. Initial reslice + camera fit ────────────────────────────────
+    # Patch on_data_ready to fit the camera on the first delivery.
+    scene_mgr = controller._render_manager._scenes[scene.id]
+    gfx_vis = scene_mgr.get_visual(visual.id)
+    _camera_fitted: set[str] = set()
+
+    _orig_3d = gfx_vis.on_data_ready
+
+    def _patched_3d(batch):
+        _orig_3d(batch)
+        if "3d" not in _camera_fitted:
+            _camera_fitted.add("3d")
+            controller.look_at_visual(
+                visual.id, view_direction=(-1, -1, -1), up=(0, 0, 1)
+            )
+
+    gfx_vis.on_data_ready = _patched_3d
+
+    _orig_2d = gfx_vis.on_data_ready_2d
+
+    def _patched_2d(batch):
+        _orig_2d(batch)
+        if "2d" not in _camera_fitted:
+            _camera_fitted.add("2d")
+            controller.look_at_visual(
+                visual.id, view_direction=(0, 0, -1), up=(0, 1, 0)
+            )
+
+    gfx_vis.on_data_ready_2d = _patched_2d
 
     QTimer.singleShot(0, controller.reslice_all)
 
 
@@ -24,6 +24,7 @@
     QCheckBox,
     QDoubleSpinBox,
     QFrame,
+    QGroupBox,
     QHBoxLayout,
     QLabel,
     QMainWindow,
@@ -126,47 +127,6 @@ def _setup_debug_logging(spec: str) -> None:
 # ---------------------------------------------------------------------------
 
 
-def _make_box_wireframe(
-    box_min: np.ndarray, box_max: np.ndarray, color: str
-) -> gfx.Line:
-    """Build an AABB wireframe as disconnected edge pairs."""
-    x0, y0, z0 = box_min
-    x1, y1, z1 = box_max
-    positions = np.array(
-        [
-            [x0, y0, z0],
-            [x1, y0, z0],
-            [x1, y0, z0],
-            [x1, y1, z0],
-            [x1, y1, z0],
-            [x0, y1, z0],
-            [x0, y1, z0],
-            [x0, y0, z0],
-            [x0, y0, z1],
-            [x1, y0, z1],
-            [x1, y0, z1],
-            [x1, y1, z1],
-            [x1, y1, z1],
-            [x0, y1, z1],
-            [x0, y1, z1],
-            [x0, y0, z1],
-            [x0, y0, z0],
-            [x0, y0, z1],
-            [x1, y0, z0],
-            [x1, y0, z1],
-            [x1, y1, z0],
-            [x1, y1, z1],
-            [x0, y1, z0],
-            [x0, y1, z1],
-        ],
-        dtype=np.float32,
-    )
-    return gfx.Line(
-        gfx.Geometry(positions=positions),
-        gfx.LineSegmentMaterial(color=color, thickness=1.0),
-    )
-
-
 def _make_frustum_wireframe(corners: np.ndarray, color: str = "#00cc44") -> gfx.Line:
     """Build a wireframe from 8 frustum corners (shape 2x4x3)."""
     edge_indices = [
@@ -269,21 +229,10 @@ def __init__(self, data_store: MultiscaleZarrDataStore) -> None:
             gpu_budget_bytes=GPU_BUDGET_3D,
             threshold=0.2,
         )
+        self._visual_2d.aabb.color = AABB_COLOR
+        self._visual_3d.aabb.color = AABB_COLOR
         self._canvas_widget_3d = self._controller.add_canvas(self._scene_3d.id)
 
-        # Add AABB wireframe to the 3D scene.
-        gfx_visual_3d = self._get_gfx_visual_3d()
-        d, h, w = gfx_visual_3d._volume_geometry.base_layout.volume_shape
-        pad = 1.0
-        gfx_scene_3d = self._controller._render_manager.get_scene(self._scene_3d.id)
-        gfx_scene_3d.add(
-            _make_box_wireframe(
-                np.array([-0.5 - pad, -0.5 - pad, -0.5 - pad]),
-                np.array([w - 0.5 + pad, h - 0.5 + pad, d - 0.5 + pad]),
-                AABB_COLOR,
-            )
-        )
-
         self._setup_ui()
 
     # ── Render-layer visual accessors ─────────────────────────────────
@@ -417,6 +366,25 @@ def _setup_ui(self) -> None:
         self._lod_bias_sb.setValue(LOD_BIAS)
         pl.addWidget(self._lod_bias_sb)
 
+        pl.addWidget(_make_separator())
+
+        # ── Shared: bounding box ─────────────────────────────────────
+        aabb_group = QGroupBox("Bounding box")
+        aabb_layout = QVBoxLayout(aabb_group)
+        self._show_aabb_cb = QCheckBox("Show bounding box")
+        self._show_aabb_cb.setChecked(False)
+        self._show_aabb_cb.toggled.connect(self._on_show_aabb_toggled)
+        aabb_layout.addWidget(self._show_aabb_cb)
+        lw_row = QHBoxLayout()
+        lw_row.addWidget(QLabel("Line width (px):"))
+        self._aabb_lw_sb = QSpinBox()
+        self._aabb_lw_sb.setRange(1, 20)
+        self._aabb_lw_sb.setValue(int(self._visual_3d.aabb.line_width))
+        self._aabb_lw_sb.valueChanged.connect(self._on_aabb_line_width_changed)
+        lw_row.addWidget(self._aabb_lw_sb)
+        aabb_layout.addLayout(lw_row)
+        pl.addWidget(aabb_group)
+
         pl.addStretch()
 
         self._status_label = QLabel("Auto-redraw active — move camera to update")
@@ -502,6 +470,14 @@ def _on_show_frustum_toggled(self, checked: bool) -> None:
         if self._frustum_line is not None:
             self._frustum_line.visible = checked
 
+    def _on_show_aabb_toggled(self, checked: bool) -> None:
+        self._visual_2d.aabb.enabled = checked
+        self._visual_3d.aabb.enabled = checked
+
+    def _on_aabb_line_width_changed(self, value: int) -> None:
+        self._visual_2d.aabb.line_width = float(value)
+        self._visual_3d.aabb.line_width = float(value)
+
     def _on_far_plane_changed(self, value: float) -> None:
         self._controller.set_camera_depth_range(
             scene_id=self._scene_3d.id,