Cleaned imports, more config options

nerfstudio/configs/method_configs.py

@@ -287,7 +287,7 @@ method_configs["phototourism"] = ExperimentConfig(
     vis="viewer",
 )
 
-method_configs["dreamfusion"] = Config(
+method_configs["dreamfusion"] = ExperimentConfig(
     method_name="dreamfusion",
     trainer=TrainerConfig(
         steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -18,64 +18,24 @@ Data manager for dreamfusion
 
 from __future__ import annotations
 
-import random
-from abc import abstractmethod
-from copy import deepcopy
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import Dict, List, Tuple, Type, Union
 
-import numpy as np
-import numpy.typing as npt
 import torch
-import torch.nn.functional as F
-import tyro
-from PIL import Image
 from rich.progress import Console
-from torch import nn
 from torch.nn import Parameter
-from torch.utils.data import Dataset
-from torch.utils.data.distributed import DistributedSampler
-from torchtyping import TensorType
 from typing_extensions import Literal
 
-import nerfstudio.utils.profiler as profiler
-from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.cameras.rays import RayBundle
-from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.configs.config_utils import to_immutable_dict
 from nerfstudio.data.datamanagers.base_datamanager import (
     DataManager,
     VanillaDataManager,
     VanillaDataManagerConfig,
 )
-from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
-from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
-from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
-from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
-from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
-    InstantNGPDataParserConfig,
-)
-from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
-from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig
-from nerfstudio.data.dataparsers.phototourism_dataparser import (
-    PhototourismDataParserConfig,
-)
-from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig
 from nerfstudio.data.datasets.base_dataset import InputDataset
-from nerfstudio.data.pixel_samplers import PixelSampler
 from nerfstudio.data.scene_box import SceneBox
-from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path
-from nerfstudio.data.utils.dataloaders import (
-    CacheDataloader,
-    FixedIndicesEvalDataloader,
-    RandIndicesEvalDataloader,
-)
-from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
-from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
-from nerfstudio.engine.optimizers import AdamOptimizerConfig
-from nerfstudio.model_components.ray_generators import RayGenerator
-from nerfstudio.utils.misc import IterableWrapper
 
 CONSOLE = Console(width=120)
 
@@ -104,9 +64,9 @@ def random_train_pose(
     device,
     radius_mean=1.0,
     radius_std=0.1,
-    central_rotation_range=[0, 360],
-    vertical_rotation_range=[-90, 10],
-    focal_range=[0.75, 1.35],
+    central_rotation_range=(0, 360),
+    vertical_rotation_range=(-90, 10),
+    focal_range=(0.75, 1.35),
     jitter_std=0.01,
 ):
     """generate random poses from an orbit camera

nerfstudio/engine/trainer.py

@@ -149,9 +149,10 @@ class Trainer:
         """Train the model."""
         assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs"
 
-        self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform(
-            self.base_dir / "dataparser_transforms.json"
-        )
+        if not self.pipeline.generative:
+            self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform(
+                self.base_dir / "dataparser_transforms.json"
+            )
 
         self._init_viewer_state()
         with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):
@@ -333,8 +334,10 @@ class Trainer:
         self.optimizers.zero_grad_all()
         cpu_or_cuda_str = self.device.split(":")[0]
 
-        if hasattr(self.pipeline, 'custom_step'):
-            _, loss_dict, metrics_dict = self.pipeline.custom_step(step=step, grad_scaler=self.grad_scaler, optimizers=self.optimizers)
+        if hasattr(self.pipeline, "custom_step"):
+            _, loss_dict, metrics_dict = self.pipeline.custom_step(
+                step=step, grad_scaler=self.grad_scaler, optimizers=self.optimizers
+            )
             loss = functools.reduce(torch.add, loss_dict.values())
         else:
             with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision):

nerfstudio/pipelines/base_pipeline.py

@@ -94,6 +94,7 @@ class Pipeline(nn.Module):
 
     datamanager: DataManager
     _model: Model
+    generative: bool
 
     @property
     def model(self):
@@ -221,6 +222,7 @@ class VanillaPipeline(Pipeline):
         self.datamanager: VanillaDataManager = config.datamanager.setup(
             device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank
         )
+        self.generative = False
         self.datamanager.to(device)
         # TODO(ethan): get rid of scene_bounds from the model
         assert self.datamanager.train_dataset is not None, "Missing input dataset"
@@ -229,7 +231,7 @@ class VanillaPipeline(Pipeline):
             scene_box=self.datamanager.train_dataset.scene_box,
             num_train_data=len(self.datamanager.train_dataset),
             metadata=self.datamanager.train_dataset.metadata,
-            device=device
+            device=device,
         )
         self.model.to(device)
 

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -29,6 +29,14 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     """enables location based prompting"""
     opacity_penalty: bool = True
     """enables penalty to encourage transparent scenes"""
+    target_transmittance: float = 0.25
+    """target transmittance for opacity penalty"""
+    side_prompt: str = ", side view"
+    """appended to prompt for side view"""
+    front_prompt: str = ", front view"
+    """appended to prompt for front view"""
+    back_prompt: str = ", back view"
+    """appended to prompt for back view"""
 
 
 class DreamfusionPipeline(VanillaPipeline):
@@ -41,12 +49,13 @@ class DreamfusionPipeline(VanillaPipeline):
         local_rank: int = 0,
     ):
         super().__init__(config, device, test_mode, world_size, local_rank)
+        self.generative = True
         self.sd = StableDiffusion(device)
         self.base_text_embedding = self.sd.get_text_embeds(config.prompt, "")
         if config.location_based_prompting:
-            self.front_text_embedding = self.sd.get_text_embeds(f"{config.prompt}, front view", "")
-            self.side_text_embedding = self.sd.get_text_embeds(f"{config.prompt}, side view", "")
-            self.back_text_embedding = self.sd.get_text_embeds(f"{config.prompt}, back view", "")
+            self.front_text_embedding = self.sd.get_text_embeds(f"{config.prompt}{config.front_prompt}", "")
+            self.side_text_embedding = self.sd.get_text_embeds(f"{config.prompt}{config.side_prompt}", "")
+            self.back_text_embedding = self.sd.get_text_embeds(f"{config.prompt}{config.back_prompt}", "")
 
     @profiler.time_function
     def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers):