Remove dataparser

nerfstudio/configs/experiment_config.py

@@ -72,7 +72,9 @@ class ExperimentConfig(InstantiateConfig):
     vis: Literal["viewer", "wandb", "tensorboard", "viewer+wandb", "viewer+tensorboard"] = "wandb"
     """Which visualizer to use."""
     data: Optional[Path] = None
-    """Alias for --pipeline.datamanager.dataparser.data"""
+    """Alias for --pipeline.datamanager.data"""
+    prompt: Optional[str] = None
+    """Alias for --pipeline.model.prompt"""
     relative_model_dir: Path = Path("nerfstudio_models/")
     """Relative path to save all checkpoints."""
 
@@ -96,7 +98,7 @@ class ExperimentConfig(InstantiateConfig):
     def set_experiment_name(self) -> None:
         """Dynamically set the experiment name"""
         if self.experiment_name is None:
-            self.experiment_name = str(self.pipeline.datamanager.dataparser.data).replace("../", "").replace("/", "-")
+            self.experiment_name = str(self.pipeline.datamanager.data).replace("../", "").replace("/", "-")
 
     def get_base_dir(self) -> Path:
         """Retrieve the base directory to set relative paths"""

nerfstudio/configs/method_configs.py

@@ -338,6 +338,7 @@ method_configs["phototourism"] = TrainerConfig(
 
 method_configs["dreamfusion"] = TrainerConfig(
     method_name="dreamfusion",
+    experiment_name="",
     steps_per_eval_batch=50,
     steps_per_eval_image=50,
     steps_per_save=200,
@@ -346,9 +347,6 @@ method_configs["dreamfusion"] = TrainerConfig(
     pipeline=VanillaPipelineConfig(
         generative=True,
         datamanager=DreamFusionDataManagerConfig(
-            dataparser=NerfstudioDataParserConfig(),
-            train_num_rays_per_batch=4096,
-            eval_num_rays_per_batch=4096,
             horizontal_rotation_warmup=4000,
         ),
         model=DreamFusionModelConfig(

nerfstudio/data/datamanagers/base_datamanager.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 
 from abc import abstractmethod
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import torch
@@ -83,6 +84,21 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[  # Omit prefixes of
 same as the vanilla union, but results in shorter subcommand names."""
 
 
+@dataclass
+class DataManagerConfig(InstantiateConfig):
+    """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
+    After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked
+    train/eval data at each iteration
+    """
+
+    _target: Type = field(default_factory=lambda: DataManager)
+    """Target class to instantiate."""
+    data: Optional[Path] = None
+    """Source of data, may not be used by all models."""
+    camera_optimizer: Optional[CameraOptimizerConfig] = None
+    """Specifies the camera pose optimizer used during training. Helpful if poses are noisy."""
+
+
 class DataManager(nn.Module):
     """Generic data manager's abstract class
 
@@ -201,12 +217,10 @@ class DataManager(nn.Module):
         """Sets up the data manager for training.
 
         Here you will define any subclass specific object attributes from the attribute"""
-        raise NotImplementedError
 
     @abstractmethod
     def setup_eval(self):
         """Sets up the data manager for evaluation"""
-        raise NotImplementedError
 
     @abstractmethod
     def next_train(self, step: int) -> Tuple:
@@ -229,6 +243,20 @@ class DataManager(nn.Module):
         """Returns the next eval image."""
         raise NotImplementedError
 
+    @abstractmethod
+    def get_train_rays_per_batch(self) -> int:
+        """Returns the number of rays per batch for training."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_eval_rays_per_batch(self) -> int:
+        """Returns the number of rays per batch for evaluation."""
+        raise NotImplementedError
+
+    def get_datapath(self) -> Optional[Path]:  # pylint:disable=no-self-use
+        """Returns the path to the data. This is used to determine where to save camera paths."""
+        return None
+
     def get_training_callbacks(  # pylint:disable=no-self-use
         self, training_callback_attributes: TrainingCallbackAttributes  # pylint: disable=unused-argument
     ) -> List[TrainingCallback]:
@@ -246,7 +274,7 @@ class DataManager(nn.Module):
 
 
 @dataclass
-class VanillaDataManagerConfig(InstantiateConfig):
+class VanillaDataManagerConfig(DataManagerConfig):
     """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
     After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked
     train/eval data at each iteration
@@ -319,7 +347,12 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.sampler = None
         self.test_mode = test_mode
         self.test_split = "test" if test_mode in ["test", "inference"] else "val"
-        self.dataparser = self.config.dataparser.setup()
+        self.dataparser = self.config.dataparser
+        if self.config.data is not None:
+            self.config.dataparser.data = Path(self.config.data)
+        else:
+            self.config.data = self.config.dataparser.data
+        self.dataparser = self.dataparser.setup()
         self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train")
 
         self.train_dataset = self.create_train_dataset()
@@ -438,6 +471,15 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             return image_idx, camera_ray_bundle, batch
         raise ValueError("No more eval images")
 
+    def get_train_rays_per_batch(self) -> int:
+        return self.config.train_num_rays_per_batch
+
+    def get_eval_rays_per_batch(self) -> int:
+        return self.config.eval_num_rays_per_batch
+
+    def get_datapath(self) -> Path:
+        return self.config.dataparser.data
+
     def get_param_groups(self) -> Dict[str, List[Parameter]]:  # pylint: disable=no-self-use
         """Get the param groups for the data manager.
         Returns:

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -18,6 +18,7 @@ Data manager for dreamfusion
 
 from __future__ import annotations
 
+import sys
 from dataclasses import dataclass, field
 from typing import Dict, List, Tuple, Type, Union
 
@@ -29,11 +30,7 @@ from typing_extensions import Literal
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.configs.config_utils import to_immutable_dict
-from nerfstudio.data.datamanagers.base_datamanager import (
-    DataManager,
-    VanillaDataManager,
-    VanillaDataManagerConfig,
-)
+from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManagerConfig
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.scene_box import SceneBox
 
@@ -146,11 +143,8 @@ def random_train_pose(
 
 
 @dataclass
-class DreamFusionDataManagerConfig(VanillaDataManagerConfig):
-    """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
-    After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked
-    train/eval data at each iteration
-    """
+class DreamFusionDataManagerConfig(DataManagerConfig):
+    """Configuration for data manager that does not load from a dataset. Instead, it generates random poses."""
 
     _target: Type = field(default_factory=lambda: DreamFusionDataManager)
     train_resolution: int = 64
@@ -179,7 +173,7 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig):
     """How many steps until the full horizontal rotation range is used"""
 
 
-class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-method
+class DreamFusionDataManager(DataManager):  # pylint: disable=abstract-method
     """Basic stored data manager implementation.
 
     This is pretty much a port over from our old dataloading utilities, and is a little jank
@@ -211,7 +205,10 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
         self.sampler = None
         self.test_mode = test_mode
         self.test_split = "test" if test_mode in ["test", "inference"] else "val"
-        self.dataparser = self.config.dataparser.setup()
+
+        if self.config.data is not None:
+            CONSOLE.print("[red] --data should not be used with the DreamFusionDataManager[/red]")
+            sys.exit(1)
 
         cameras, _, _ = random_train_pose(
             self.config.num_eval_angles,
@@ -266,7 +263,11 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
         # camera_idx = torch.randint(0, self.eval_cameras.shape[0], [1], dtype=torch.long, device=self.device)
         # ray_bundle = self.eval_cameras.generate_rays(camera_idx).flatten()
 
-        return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation, "initialization": True}
+        return ray_bundle, {
+            "vertical": vertical_rotation,
+            "central": central_rotation,
+            "initialization": True,
+        }
 
     def next_eval(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the eval dataloader."""
@@ -289,6 +290,12 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
 
         return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation}
 
+    def get_train_rays_per_batch(self) -> int:
+        return self.config.train_resolution**2
+
+    def get_eval_rays_per_batch(self) -> int:
+        return self.config.eval_resolution**2
+
     def get_param_groups(self) -> Dict[str, List[Parameter]]:  # pylint: disable=no-self-use
         """Get the param groups for the data manager.
         Returns:

nerfstudio/engine/trainer.py

@@ -117,24 +117,12 @@ class Trainer:
         # optimizers
         self.grad_scaler = GradScaler(enabled=self.mixed_precision)
 
+        self.viewer_state = None
+
         self.base_dir = config.get_base_dir()
         # directory to save checkpoints
         self.checkpoint_dir = config.get_checkpoint_dir()
         CONSOLE.log(f"Saving checkpoints to: {self.checkpoint_dir}")
-        # set up viewer if enabled
-        viewer_log_path = self.base_dir / config.viewer.relative_log_filename
-        self.viewer_state, banner_messages = None, None
-        if self.config.is_viewer_enabled() and local_rank == 0:
-            self.viewer_state, banner_messages = viewer_utils.setup_viewer(
-                config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data
-            )
-        self._check_viewer_warnings()
-        # set up writers/profilers if enabled
-        writer_log_path = self.base_dir / config.logging.relative_log_dir
-        writer.setup_event_writer(config.is_wandb_enabled(), config.is_tensorboard_enabled(), log_dir=writer_log_path)
-        writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages)
-        writer.put_config(name="config", config_dict=dataclasses.asdict(config), step=0)
-        profiler.setup_profiler(config.logging)
 
     def setup(self, test_mode: Literal["test", "val", "inference"] = "val"):
         """Setup the Trainer by calling other setup functions.
@@ -164,6 +152,28 @@ class Trainer:
             )
         )
 
+        # set up viewer if enabled
+        viewer_log_path = self.base_dir / self.config.viewer.relative_log_filename
+        self.viewer_state, banner_messages = None, None
+        if self.config.is_viewer_enabled() and self.local_rank == 0:
+            datapath = self.pipeline.datamanager.get_datapath()
+            if datapath is None:
+                datapath = self.base_dir
+            self.viewer_state, banner_messages = viewer_utils.setup_viewer(
+                self.config.viewer, log_filename=viewer_log_path, datapath=datapath
+            )
+        self._check_viewer_warnings()
+        # set up writers/profilers if enabled
+        writer_log_path = self.base_dir / self.config.logging.relative_log_dir
+        writer.setup_event_writer(
+            self.config.is_wandb_enabled(), self.config.is_tensorboard_enabled(), log_dir=writer_log_path
+        )
+        writer.setup_local_writer(
+            self.config.logging, max_iter=self.config.max_num_iterations, banner_messages=banner_messages
+        )
+        writer.put_config(name="config", config_dict=dataclasses.asdict(self.config), step=0)
+        profiler.setup_profiler(self.config.logging)
+
     def setup_optimizers(self) -> Optimizers:
         """Helper to set up the optimizers
 
@@ -171,9 +181,9 @@ class Trainer:
             The optimizers object given the trainer config.
         """
         optimizer_config = self.config.optimizers.copy()
-        camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer
         param_groups = self.pipeline.get_param_groups()
-        if camera_optimizer_config.mode != "off":
+        camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer
+        if camera_optimizer_config is not None and camera_optimizer_config.mode != "off":
             assert camera_optimizer_config.param_group not in optimizer_config
             optimizer_config[camera_optimizer_config.param_group] = {
                 "optimizer": camera_optimizer_config.optimizer,
@@ -215,7 +225,7 @@ class Trainer:
                 if step > 1:
                     writer.put_time(
                         name=EventName.TRAIN_RAYS_PER_SEC,
-                        duration=self.config.pipeline.datamanager.train_num_rays_per_batch / train_t.duration,
+                        duration=self.pipeline.datamanager.get_train_rays_per_batch() / train_t.duration,
                         step=step,
                         avg_over_steps=True,
                     )
@@ -249,7 +259,7 @@ class Trainer:
 
     @check_main_thread
     def _always_render(self, step):
-        if self.config.is_viewer_enabled():
+        if self.viewer_state is not None:
             while True:
                 self.viewer_state.vis["renderingState/isTraining"].write(False)
                 self._update_viewer_state(step)
@@ -289,7 +299,7 @@ class Trainer:
         """
         assert self.viewer_state is not None
         with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _:
-            num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
+            num_rays_per_batch = self.pipeline.datamanager.get_train_rays_per_batch()
             try:
                 self.viewer_state.update_scene(self, step, self.pipeline.model, num_rays_per_batch)
             except RuntimeError:
@@ -308,7 +318,7 @@ class Trainer:
             vis_t: timer object carrying time to execute visualization step
             step: current step
         """
-        train_num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
+        train_num_rays_per_batch = self.pipeline.datamanager.get_train_rays_per_batch()
         writer.put_time(
             name=EventName.TRAIN_RAYS_PER_SEC,
             duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration),

nerfstudio/models/base_model.py

@@ -50,6 +50,8 @@ class ModelConfig(InstantiateConfig):
     """parameters to instantiate density field with"""
     eval_num_rays_per_chunk: int = 4096
     """specifies number of rays per chunk during eval"""
+    prompt: Optional[str] = None
+    """A prompt to be used in text to NeRF models"""
 
 
 class Model(nn.Module):

nerfstudio/pipelines/base_pipeline.py

@@ -41,6 +41,7 @@ from typing_extensions import Literal
 from nerfstudio.configs import base_config as cfg
 from nerfstudio.data.datamanagers.base_datamanager import (
     DataManager,
+    DataManagerConfig,
     VanillaDataManager,
     VanillaDataManagerConfig,
 )
@@ -187,7 +188,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig):
 
     _target: Type = field(default_factory=lambda: VanillaPipeline)
     """target class to instantiate"""
-    datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig()
+    datamanager: DataManagerConfig = VanillaDataManagerConfig()
     """specifies the datamanager config"""
     model: ModelConfig = ModelConfig()
     """specifies the model config"""
@@ -264,15 +265,16 @@ class VanillaPipeline(Pipeline):
         model_outputs = self.model(ray_bundle)
         metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
 
-        camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group
-        if camera_opt_param_group in self.datamanager.get_param_groups():
-            # Report the camera optimization metrics
-            metrics_dict["camera_opt_translation"] = (
-                self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()
-            )
-            metrics_dict["camera_opt_rotation"] = (
-                self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()
-            )
+        if self.config.datamanager.camera_optimizer is not None:
+            camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group
+            if camera_opt_param_group in self.datamanager.get_param_groups():
+                # Report the camera optimization metrics
+                metrics_dict["camera_opt_translation"] = (
+                    self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()
+                )
+                metrics_dict["camera_opt_rotation"] = (
+                    self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()
+                )
 
         loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
 

scripts/train.py

@@ -219,8 +219,12 @@ def main(config: TrainerConfig) -> None:
 
     config.set_timestamp()
     if config.data:
-        CONSOLE.log("Using --data alias for --data.pipeline.datamanager.dataparser.data")
-        config.pipeline.datamanager.dataparser.data = config.data
+        CONSOLE.log("Using --data alias for --data.pipeline.datamanager.data")
+        config.pipeline.datamanager.data = config.data
+
+    if config.prompt:
+        CONSOLE.log("Using --prompt alias for --data.pipeline.model.prompt")
+        config.pipeline.model.prompt = config.prompt
 
     if config.load_config:
         CONSOLE.log(f"Loading pre-set config from: {config.load_config}")

scripts/viewer/run_viewer.py

@@ -76,16 +76,16 @@ class RunViewer:
         )
         while True:
             viewer_state.vis["renderingState/isTraining"].write(False)
-            self._update_viewer_state(viewer_state, config, pipeline)
+            self._update_viewer_state(viewer_state, pipeline)
 
-    def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, config: TrainerConfig, pipeline: 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 = config.pipeline.datamanager.train_num_rays_per_batch
+        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)