diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 2ee48ca75..a516ba323 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -60,7 +60,6 @@ from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig from nerfstudio.models.tensorf import TensoRFModelConfig from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig -from nerfstudio.pipelines.dreamfusion_pipeline import DreamfusionPipelineConfig from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig from nerfstudio.plugins.registry import discover_methods @@ -344,7 +343,8 @@ method_configs["dreamfusion"] = TrainerConfig( steps_per_save=200, max_num_iterations=30000, mixed_precision=True, - pipeline=DreamfusionPipelineConfig( + pipeline=VanillaPipelineConfig( + generative=True, datamanager=DreamFusionDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, @@ -352,22 +352,22 @@ method_configs["dreamfusion"] = TrainerConfig( ), model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, - distortion_loss_mult=10.0, + distortion_loss_mult=0.02, + interlevel_loss_mult=1.0, orientation_loss_mult=0.1, max_res=256, sphere_collider=True, initialize_density=False, random_background=True, - interlevel_loss_mult=1.0, - proposal_warmup=500, - proposal_update_every=5, + proposal_warmup=1000, + proposal_update_every=2, proposal_weights_anneal_max_num_iters=100, start_lambertian_training=1000, start_normals_training=500, opacity_loss_mult=0.001, + prompting_type="location_based", + guidance_scale=100, ), - interpolated_prompting=False, - guidance_scale=100, ), optimizers={ "proposal_networks": { diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py new file mode 100644 index 000000000..82f61ff7a --- /dev/null +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -0,0 +1,61 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Stable diffusion utils +""" + +def get_text_embedding(batch, prompting_type, text_embeddings, sd_device): + """""" + + base = text_embeddings["base_text_embedding"] + front = text_embeddings["front_text_embedding"] + side = text_embeddings["side_text_embedding"] + top = text_embeddings["top_text_embedding"] + back = text_embeddings["back_text_embedding"] + + assert prompting_type in ("location_based", "interpolated", "base") + + if prompting_type == "location_based": + if batch["vertical"] < 40: + text_embedding = top + elif batch["central"] > 315 or batch["central"] <= 45: + text_embedding = front + elif batch["central"] > 45 and batch["central"] <= 135: + text_embedding = side + elif batch["central"] > 135 and batch["central"] <= 225: + text_embedding = back + else: # batch["central"] > 225 and batch["central"] <= 315: + text_embedding = side + + elif prompting_type == "interpolated": + horiz = batch["central"].to(sd_device) + vert = max(batch["vertical"].to(sd_device), 0) + + if batch["central"] > 0 or batch["central"] <= 90: + text_embedding = (horiz) * side + (90 - horiz) * front + elif batch["central"] > 90 and batch["central"] <= 180: + text_embedding = (horiz - 90) * back + (180 - horiz) * side + elif batch["central"] > 180 and batch["central"] <= 270: + text_embedding = (horiz - 180) * side + (270 - horiz) * back + else: # batch["central"] > 270 and batch["central"] <= 360: + text_embedding = (horiz - 270) * front + (360 - horiz) * side + + text_embedding = text_embedding / 90.0 + text_embedding = (vert * text_embedding + (90 - vert) * top) / 90.0 + + elif prompting_type == "base": + text_embedding = base + + return text_embedding diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 49baefd98..f55795ab5 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -510,7 +510,7 @@ class VolumetricSampler(Sampler): class ProposalNetworkSampler(Sampler): """Sampler that uses a proposal network to generate samples. - + Args: num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step. num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model. @@ -518,6 +518,7 @@ class ProposalNetworkSampler(Sampler): single_jitter: Use a same random jitter for all samples along a ray. update_sched: A function that takes the iteration number of steps between updates. initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set. + pdf_histogram_padding: Padding to add to the histogram of the pdf. """ def __init__( @@ -528,6 +529,7 @@ class ProposalNetworkSampler(Sampler): single_jitter: bool = False, update_sched: Callable = lambda x: 1, initial_sampler: Optional[Sampler] = None, + pdf_histogram_padding: float = 0.01, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray @@ -542,7 +544,9 @@ class ProposalNetworkSampler(Sampler): self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) else: self.initial_sampler = initial_sampler - self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) + self.pdf_sampler = PDFSampler( + include_original=False, single_jitter=single_jitter, histogram_padding=pdf_histogram_padding + ) self._anneal = 1.0 self._steps_since_update = 0 diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index c2a716946..6eba9dc96 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -19,11 +19,12 @@ DreamFusion implementation. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Tuple, Type +from typing import Dict, List, Tuple, Type, Optional import numpy as np import torch from torch.nn import Parameter +from torch.cuda.amp.grad_scaler import GradScaler from nerfstudio.cameras.rays import RayBundle from nerfstudio.engine.callbacks import ( @@ -31,6 +32,8 @@ from nerfstudio.engine.callbacks import ( TrainingCallbackAttributes, TrainingCallbackLocation, ) +from nerfstudio.generative.stable_diffusion import StableDiffusion +from nerfstudio.generative.stable_diffusion_utils import get_text_embedding from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.fields.dreamfusion_field import DreamFusionField @@ -128,6 +131,25 @@ class DreamFusionModelConfig(ModelConfig): max_res: int = 256 """Maximum resolution of the density field.""" + location_based_prompting: bool = True + """enables location based prompting""" + interpolated_prompting: bool = False + """enables interpolated location prompting""" + prompting_type: str = "location_based" + """choose between location_based, interpolated, base""" + top_prompt: str = ", overhead view" + """appended to prompt for overhead view""" + 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""" + guidance_scale: float = 100 + """guidance scale for sds loss""" + stablediffusion_device: Optional[str] = None + """device for stable diffusion""" + sd_version: str = '1-5' class DreamFusionModel(Model): """DreamFusionModel Model @@ -144,18 +166,44 @@ class DreamFusionModel(Model): **kwargs, ) -> None: self.prompt = config.prompt + self.cur_prompt = config.prompt + self.sd_version = config.sd_version self.initialize_density = config.initialize_density self.train_normals = False self.train_shaded = False self.random_background = config.random_background self.density_strength = 1.0 self.target_transmittance = config.target_transmittance_start + self.grad_scaler = kwargs['grad_scaler'] + + self.prompting_type = config.prompting_type + self.guidance_scale = config.guidance_scale + self.top_prompt = config.top_prompt + self.side_prompt = config.side_prompt + self.back_prompt = config.back_prompt + self.front_prompt = config.front_prompt + + self.sd_device = ( + torch.device(kwargs['device']) + if config.stablediffusion_device is None + else torch.device(config.stablediffusion_device) + ) + super().__init__(config=config, **kwargs) def populate_modules(self): """Set the fields and modules""" super().populate_modules() + self.sd = StableDiffusion(self.sd_device, version=self.sd_version) + self.text_embeddings = { + "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""), + "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""), + "side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""), + "back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""), + "base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "") + } + # setting up fields self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res) @@ -186,6 +234,7 @@ class DreamFusionModel(Model): single_jitter=self.config.use_single_jitter, update_sched=update_schedule, initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter), + pdf_histogram_padding=0.001, ) # renderers @@ -350,7 +399,7 @@ class DreamFusionModel(Model): shading_weight = 0.0 shaded, shaded_albedo = self.shader_lambertian( - rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True + rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False ) outputs["normals"] = self.shader_normals(normals, weights=accum_mask) @@ -412,6 +461,33 @@ class DreamFusionModel(Model): if self.config.opacity_penalty: loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean() + if self.prompt != self.cur_prompt: + self.cur_prompt = self.prompt + self.text_embeddings = { + "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""), + "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""), + "side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""), + "back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""), + "base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "") + } + + text_embedding = get_text_embedding(batch, self.prompting_type, self.text_embeddings, self.sd_device) + + train_output = ( + outputs["train_output"].view(1, int(outputs["train_output"].shape[0] ** 0.5), int(outputs["train_output"].shape[0] ** 0.5), 3) + .permute(0, 3, 1, 2) + ) + + with torch.autocast(device_type="cuda", dtype=torch.float16): + sds_loss = self.sd.sds_loss( + text_embedding.to(self.sd_device), + train_output.to(self.sd_device), + guidance_scale=int(self.guidance_scale), + grad_scaler=self.grad_scaler, + ) + + loss_dict["sds_loss"] = sds_loss.to(self.device) + if self.training: loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( outputs["weights_list"], outputs["ray_samples_list"] diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 3b7c45ecb..ac1b73dd6 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -191,6 +191,8 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """specifies the datamanager config""" model: ModelConfig = ModelConfig() """specifies the model config""" + generative: bool = False + """specifies whether the pipeline is for generative models""" class VanillaPipeline(Pipeline): @@ -225,7 +227,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.generative = config.generative 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" @@ -235,6 +237,7 @@ class VanillaPipeline(Pipeline): num_train_data=len(self.datamanager.train_dataset), metadata=self.datamanager.train_dataset.metadata, device=device, + grad_scaler=grad_scaler ) self.model.to(device) diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py deleted file mode 100644 index 63f86816b..000000000 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Dreamfusion Pipeline and trainer""" - -from dataclasses import dataclass, field -from typing import Optional, Type - -import torch -from torch.cuda.amp.grad_scaler import GradScaler -from typing_extensions import Literal - -from nerfstudio.data.datamanagers.dreamfusion_datamanager import ( - DreamFusionDataManagerConfig, -) -from nerfstudio.generative.stable_diffusion import StableDiffusion -from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig - - -@dataclass -class DreamfusionPipelineConfig(VanillaPipelineConfig): - """Configuration for pipeline instantiation""" - - _target: Type = field(default_factory=lambda: DreamfusionPipeline) - """target class to instantiate""" - datamanager: DreamFusionDataManagerConfig = DreamFusionDataManagerConfig() - """specifies the datamanager config""" - # prompt: str = "A high-quality photo of a tree frog on a stump" - # """prompt for stable dreamfusion""" - - location_based_prompting: bool = True - """enables location based prompting""" - interpolated_prompting: bool = False - """enables interpolated location prompting""" - top_prompt: str = ", overhead view" - """appended to prompt for overhead view""" - 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""" - guidance_scale: float = 100 - """guidance scale for sds loss""" - sd_version: str = "1-5" - """Stable Diffusion version for dreamfusion""" - - stablediffusion_device: Optional[str] = None - - -class DreamfusionPipeline(VanillaPipeline): - """Dreamfusion pipeline""" - - config: DreamfusionPipelineConfig - - def __init__( - self, - config: DreamfusionPipelineConfig, - device: str, - test_mode: Literal["test", "val", "inference"] = "val", - world_size: int = 1, - local_rank: int = 0, - grad_scaler: Optional[GradScaler] = None, - ): - super().__init__(config, device, test_mode, world_size, local_rank, grad_scaler) - self.generative = True - self.cur_prompt = self.model.prompt - self.grad_scaler = grad_scaler - self.sd_device = ( - torch.device(device) - if self.config.stablediffusion_device is None - else torch.device(self.config.stablediffusion_device) - ) - self.location_based_prompting = config.location_based_prompting - self.top_prompt = config.top_prompt - self.side_prompt = config.side_prompt - self.back_prompt = config.back_prompt - self.front_prompt = config.front_prompt - self.sd = StableDiffusion(self.sd_device, version=config.sd_version) - - if config.location_based_prompting: - self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "") - self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "") - self.side_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", "") - self.back_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", "") - else: - self.base_text_embedding = self.sd.get_text_embeds(self.cur_prompt, "") - - def get_train_loss_dict(self, step: int): - - if self.world_size > 1 and step: - assert self.datamanager.train_sampler is not None - self.datamanager.train_sampler.set_epoch(step) - - if self.model.prompt != self.cur_prompt: - self.cur_prompt = self.model.prompt - if self.location_based_prompting: - self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "") - self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "") - self.side_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", "") - self.back_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", "") - else: - self.base_text_embedding = self.sd.get_text_embeds(self.cur_prompt, "") - - ray_bundle, batch = self.datamanager.next_train(step) - model_outputs = self.model(ray_bundle) - metrics_dict = self.model.get_metrics_dict(model_outputs, batch) - loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) - - train_output = ( - model_outputs["train_output"] - .view(1, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3) - .permute(0, 3, 1, 2) - ) - - if self.config.location_based_prompting: - if self.config.interpolated_prompting: - horiz = batch["central"].to(self.sd_device) - vert = max(batch["vertical"].to(self.sd_device), 0) - - if batch["central"] > 0 or batch["central"] <= 90: - text_embedding = ( - (horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding - ) / 90.0 - elif batch["central"] > 90 and batch["central"] <= 180: - text_embedding = ( - (horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding - ) / 90.0 - elif batch["central"] > 180 and batch["central"] <= 270: - text_embedding = ( - (horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding - ) / 90.0 - else: # batch["central"] > 270 and batch["central"] <= 360: - text_embedding = ( - (horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding - ) / 90.0 - - text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0 - - else: - if batch["vertical"] < 40: - text_embedding = self.top_text_embedding - elif batch["central"] > 315 or batch["central"] <= 45: - text_embedding = self.front_text_embedding - elif batch["central"] > 45 and batch["central"] <= 135: - text_embedding = self.side_text_embedding - elif batch["central"] > 135 and batch["central"] <= 225: - text_embedding = self.back_text_embedding - else: # batch["central"] > 225 and batch["central"] <= 315: - text_embedding = self.side_text_embedding - else: - text_embedding = self.base_text_embedding - - with torch.autocast(device_type="cuda", dtype=torch.float16): - sds_loss = self.sd.sds_loss( - text_embedding.to(self.sd_device), - train_output.to(self.sd_device), - guidance_scale=int(self.config.guidance_scale), - grad_scaler=self.grad_scaler, - ) - - loss_dict["sds_loss"] = sds_loss.to(self.device) - - return model_outputs, loss_dict, metrics_dict - - def forward(self): - """Not implemented since we only want the parameter saving of the nn module, but not forward()""" - raise NotImplementedError