fix viewer checkpointing (#511)

nerfactory/configs/base.py

@@ -264,6 +264,10 @@ class Config(PrintableConfig):
         self.set_experiment_name()
         return Path(f"outputs/{self.experiment_name}/{self.method_name}/{self.timestamp}")
 
+    def get_checkpoint_dir(self) -> Path:
+        """Retrieve the checkpoint directory"""
+        return Path(self.get_base_dir() / self.trainer.relative_model_dir)
+
     def print_to_terminal(self) -> None:
         """Helper to pretty print config to terminal"""
         console = Console(width=120)

nerfactory/engine/trainer.py

@@ -22,7 +22,6 @@ import functools
 import logging
 import os
 import time
-from pathlib import Path
 from typing import Any, Dict, List, Tuple
 
 import torch
@@ -102,7 +101,7 @@ class Trainer:
 
         self.base_dir = config.get_base_dir()
         # directory to save checkpoints
-        self.checkpoint_dir = Path(self.base_dir / self.config.trainer.relative_model_dir)
+        self.checkpoint_dir = config.get_checkpoint_dir()
         logging.info("Saving checkpoints to: %s", self.checkpoint_dir)
         # set up viewer if enabled
         viewer_log_path = self.base_dir / config.viewer.relative_log_filename
@@ -182,7 +181,7 @@ class Trainer:
                 self.eval_iteration(step)
 
                 if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0:
-                    self._save_checkpoint(step)
+                    self.save_checkpoint(step)
 
                 writer.write_out_storage()
 
@@ -223,7 +222,7 @@ class Trainer:
         with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _:
             num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
             try:
-                self.viewer_state.update_scene(step, self.pipeline.model, num_rays_per_batch)
+                self.viewer_state.update_scene(self, step, self.pipeline.model, num_rays_per_batch)
             except RuntimeError:
                 time.sleep(0.03)  # sleep to allow buffer to reset
                 assert self.viewer_state.vis is not None
@@ -270,7 +269,7 @@ class Trainer:
             logging.info("No checkpoints to load, training from scratch")
 
     @check_main_thread
-    def _save_checkpoint(self, step: int) -> None:
+    def save_checkpoint(self, step: int) -> None:
         """Save the model and optimizers
 
         Args:

nerfactory/viewer/server/viewer_utils.py

@@ -284,7 +284,7 @@ class ViewerState:
         # K = camera.get_intrinsics_matrix()
         # set_persp_intrinsics_matrix(self.vis, K.double().numpy())
 
-    def update_scene(self, step: int, graph: Model, num_rays_per_batch: int) -> None:
+    def update_scene(self, trainer, step: int, graph: Model, num_rays_per_batch: int) -> None:
         """updates the scene based on the graph weights
 
         Args:
@@ -298,6 +298,8 @@ class ViewerState:
         # check if we should interrupt from a button press?
         camera_path_payload = self.vis["camera_path_payload"].read()
         if camera_path_payload:
+            # save a model checkpoint
+            trainer.save_checkpoint(step)
             # write to json file
             camera_path_filename = camera_path_payload["camera_path_filename"]
             camera_path = camera_path_payload["camera_path"]