merging w generative

scripts/viewer/run_viewer.py

@@ -54,7 +54,52 @@ class RunViewer:
         config.viewer = self.viewer.as_viewer_config()
         config.viewer.num_rays_per_chunk = num_rays_per_chunk
 
-        _start_viewer(config, pipeline, step)
+        self._start_viewer(config, pipeline)
+
+    def _start_viewer(self, config: TrainerConfig, pipeline: Pipeline):
+        base_dir = config.get_base_dir()
+        viewer_log_path = base_dir / config.viewer.relative_log_filename
+
+        datapath = pipeline.datamanager.get_datapath()
+        if datapath is None:
+            datapath = base_dir
+
+        viewer_state, banner_messages = viewer_utils.setup_viewer(
+            config.viewer,
+            log_filename=viewer_log_path,
+            datapath=datapath,
+        )
+
+        # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated
+        config.logging.local_writer.enable = False
+        writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages)
+
+        assert viewer_state and pipeline.datamanager.train_dataset
+        viewer_state.init_scene(
+            dataset=pipeline.datamanager.train_dataset,
+            start_train=False,
+        )
+        while True:
+            viewer_state.vis["renderingState/isTraining"].write(False)
+            self._update_viewer_state(viewer_state, pipeline)
+
+    def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, pipeline: Pipeline):
+        """Updates the viewer state by rendering out scene with current pipeline
+        Returns the time taken to render scene.
+
+        """
+        # NOTE: step must be > 0 otherwise the rendering would not happen
+        step = 1
+        num_rays_per_batch = pipeline.datamanager.get_train_rays_per_batch()
+        with TimeWriter(writer, EventName.ITER_VIS_TIME) as _:
+            try:
+                viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch)
+            except RuntimeError:
+                time.sleep(0.03)  # sleep to allow buffer to reset
+                assert viewer_state.vis is not None
+                viewer_state.vis["renderingState/log_errors"].write(
+                    "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
+                )
 
     def save_checkpoint(self, *args, **kwargs):
         """