ngp and tensor work well now, changed shading loss for dreamfusion

scripts/viewer/run_viewer.py

@@ -54,52 +54,7 @@ class RunViewer:
         config.viewer = self.viewer.as_viewer_config()
         config.viewer.num_rays_per_chunk = num_rays_per_chunk
 
-        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."
-                )
+        _start_viewer(config, pipeline, step)
 
     def save_checkpoint(self, *args, **kwargs):
         """
@@ -109,7 +64,6 @@ class RunViewer:
 
 def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int):
     """Starts the viewer
-
     Args:
         config: Configuration of pipeline to load
         pipeline: Pipeline instance of which to load weights
@@ -117,10 +71,15 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int):
     """
     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 = ViewerState(
         config.viewer,
         log_filename=viewer_log_path,
-        datapath=pipeline.datamanager.get_datapath(),
+        datapath=datapath,
         pipeline=pipeline,
     )
     banner_messages = [f"Viewer at: {viewer_state.viewer_url}"]
@@ -147,7 +106,4 @@ def entrypoint():
 
 
 if __name__ == "__main__":
-    entrypoint()
-
-# For sphinx docs
-get_parser_fn = lambda: tyro.extras.get_parser(RunViewer)  # noqa
+    entrypoint()