diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index d0c76262d..270c7598e 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -344,7 +344,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, @@ -365,9 +366,9 @@ method_configs["dreamfusion"] = TrainerConfig( 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/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index c2a716946..c4e705664 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,23 @@ 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 class DreamFusionModel(Model): """DreamFusionModel Model @@ -144,18 +164,43 @@ class DreamFusionModel(Model): **kwargs, ) -> None: self.prompt = config.prompt + self.cur_prompt = config.prompt 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) + 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) @@ -350,7 +395,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 +457,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 index 692a92a71..88ec2feb9 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -86,7 +86,7 @@ class DreamfusionPipeline(VanillaPipeline): self.back_prompt = config.back_prompt self.front_prompt = config.front_prompt self.sd = StableDiffusion(self.sd_device) - + 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}", "") @@ -116,59 +116,59 @@ class DreamfusionPipeline(VanillaPipeline): 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) - ) + # 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 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 + # 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 + # 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 + # 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) + # 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