From 0fc4495348f297cc2d22f2339bc2c6f0939dd87e Mon Sep 17 00:00:00 2001 From: "Abhik Ahuja ahujaabhik@berkeley.edu" Date: Wed, 14 Dec 2022 05:38:17 +0000 Subject: [PATCH] Cleaned imports, more config options --- nerfstudio/configs/method_configs.py | 2 +- .../datamanagers/dreamfusion_datamanager.py | 48 ++----------------- nerfstudio/engine/trainer.py | 13 +++-- nerfstudio/pipelines/base_pipeline.py | 4 +- nerfstudio/pipelines/dreamfusion_pipeline.py | 15 ++++-- 5 files changed, 28 insertions(+), 54 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 29a708fee..1fef60978 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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 diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 9bdfc9288..7f34e5aea 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/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 diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index f283716b9..e3bd1df80 100644 --- a/nerfstudio/engine/trainer.py +++ b/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): diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 68633fa55..d63710cc0 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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) diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 3fe6cecc4..1a5845fb9 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/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):