diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 9e03323f1..c812d781f 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -414,7 +414,7 @@ method_configs["dreamplus"] = TrainerConfig( start_lambertian_training=1000, start_normals_training=500, opacity_loss_mult=0.001, - prompting_type="location_based", + positional_prompting="discrete", guidance_scale=100, ), ), diff --git a/nerfstudio/data/datamanagers/dreamfusionplus_datamanager.py b/nerfstudio/data/datamanagers/dreamfusionplus_datamanager.py index 7797b77c8..6dee16b26 100644 --- a/nerfstudio/data/datamanagers/dreamfusionplus_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusionplus_datamanager.py @@ -113,7 +113,7 @@ class DreamFusionPlusDataManagerConfig(DreamFusionDataManagerConfig): """Mean radius of the camera sphere""" use_input_transient: bool = False """Whether to allow the input image to mask out parts of the image that give poor rgb loss""" - input_distance: float = 1.5 + input_distance: float = 1.0 """Scalar multiple of radius_mean dictating how close we think this image was taken to the object in question. Large values mean this image was taken relatively far away, small values mean this was taken from relatively close to the object""" @@ -269,8 +269,11 @@ class DreamFusionPlusDataManager(DreamFusionDataManager): # pylint: disable=abs self.train_count += 1 # randomly choose a random camera view or the input view here - if np.random.random_sample() < 0.1: - return self.next_input(step) + # 0.1 - 0.5, steps 0 - 2000 + # cutoff = min(0.5, step * 0.00025) + # cutoff = 0.1 + # if np.random.random_sample() < cutoff: + # return self.next_input(step) # # TODO Reimplement when cameras are fully working # if step > 2000: diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index d7a94805e..606a96e75 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -48,6 +48,7 @@ class PositionalTextEmbeddings: stable_diffusion: StableDiffusion, positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete", ): + assert positional_prompting in ["discrete", "interpolated", "off"] self.positional_prompting = positional_prompting self.sd_device = stable_diffusion.device self.sd = stable_diffusion diff --git a/nerfstudio/models/dreamfusionplus.py b/nerfstudio/models/dreamfusionplus.py index 1c7b9eed4..5a1208f3c 100644 --- a/nerfstudio/models/dreamfusionplus.py +++ b/nerfstudio/models/dreamfusionplus.py @@ -19,7 +19,7 @@ DreamFusion implementation. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Tuple, Type, Optional +from typing import Dict, List, Tuple, Type, Optional, Literal import numpy as np import torch @@ -33,7 +33,7 @@ from nerfstudio.engine.callbacks import ( TrainingCallbackLocation, ) from nerfstudio.generative.stable_diffusion import StableDiffusion -from nerfstudio.generative.stable_diffusion_utils import get_text_embedding +from nerfstudio.generative.stable_diffusion_utils import PositionalTextEmbeddings from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.fields.dreamfusion_field import DreamFusionField @@ -78,7 +78,7 @@ class DreamFusionPlusModelConfig(DreamFusionModelConfig): """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" - init_density_strength: float = 0.1 + init_density_strength: float = 10 """Initial strength of center density""" sphere_collider: bool = True """Use spherical collider instead of box""" @@ -132,8 +132,12 @@ class DreamFusionPlusModelConfig(DreamFusionModelConfig): max_res: int = 256 """Maximum resolution of the density field.""" - prompting_type: str = "location_based" - """choose between location_based, interpolated, base""" + location_based_prompting: bool = True + """enables location based prompting""" + interpolated_prompting: bool = False + """enables interpolated location prompting""" + positional_prompting: Literal["discrete", "interpolated", "base"] = "discrete" + """ how to incorporate position into prompt""" top_prompt: str = ", overhead view" """appended to prompt for overhead view""" side_prompt: str = ", side view" @@ -146,7 +150,7 @@ class DreamFusionPlusModelConfig(DreamFusionModelConfig): """guidance scale for sds loss""" stablediffusion_device: Optional[str] = None """device for stable diffusion""" - sd_version: str = '1-5' + sd_version: str = "1-5" class DreamFusionPlusModel(DreamFusionModel): @@ -189,19 +193,21 @@ class DreamFusionPlusModel(DreamFusionModel): 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, "") - } + self.text_embeddings.update_prompt( + base_prompt=self.cur_prompt, + top_prompt=self.cur_prompt + self.top_prompt, + side_prompt=self.cur_prompt + self.side_prompt, + back_prompt=self.cur_prompt + self.back_prompt, + front_prompt=self.cur_prompt + self.front_prompt, + ) - if batch["input_image"]: + if batch["input_image"] == True: loss_dict["rgb_loss"] = self.rgb_loss(batch["image"].to(self.device), outputs["rgb"]) + loss_dict['sds_loss'] = 0 else: - text_embedding = get_text_embedding(batch, self.prompting_type, self.text_embeddings, self.sd_device) - + text_embedding = self.text_embeddings.get_text_embedding( + vertical_angle=batch["vertical"], horizonal_angle=batch["central"] + ) 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) @@ -216,12 +222,13 @@ class DreamFusionPlusModel(DreamFusionModel): ) loss_dict["sds_loss"] = sds_loss.to(self.device) + loss_dict["rgb_loss"] = 0 - if self.training: - loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( - outputs["weights_list"], outputs["ray_samples_list"] - ) - loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( - outputs["weights_list"], outputs["ray_samples_list"] - ) + if self.training: + loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) return loss_dict \ No newline at end of file