diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 30454485a..4cc56fb19 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -60,6 +60,8 @@ class DreamFusionModelConfig(ModelConfig): _target: Type = field(default_factory=lambda: DreamFusionModel) """target class to instantiate""" + prompt: str = "A high-quality photo of a tree frog on a stump" + """prompt for stable dreamfusion""" orientation_loss_mult: float = 0.0001 """Orientation loss multipier on computed normals.""" @@ -132,6 +134,7 @@ class DreamFusionModel(Model): config: DreamFusionModelConfig, **kwargs, ) -> None: + self.prompt = config.prompt self.initialize_density = config.initialize_density self.train_normals = False self.train_shaded = False @@ -339,7 +342,7 @@ class DreamFusionModel(Model): shading_weight = 0.0 shaded, shaded_albedo = self.shader_lambertian( - rgb=rgb, normals=pred_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=True ) outputs["normals"] = self.shader_normals(normals, weights=accum_mask) diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 3123bfe70..b228c7027 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -36,8 +36,8 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """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""" + # 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""" @@ -81,20 +81,27 @@ class DreamfusionPipeline(VanillaPipeline): ): 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) + if config.location_based_prompting: - self.top_text_embedding = self.sd.get_text_embeds(f"{config.prompt}{config.top_prompt}", "") - 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}", "") + 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(config.prompt, "") + self.base_text_embedding = self.sd.get_text_embeds(self.cur_prompt, "") def get_train_loss_dict(self, step: int): @@ -102,6 +109,16 @@ class DreamfusionPipeline(VanillaPipeline): 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) diff --git a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index 66b92458e..2bb6d8fd3 100644 --- a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx @@ -62,9 +62,21 @@ export function RenderControls() { const crop_center = useSelector((state) => state.renderingState.crop_center); + const generative_prompt = useSelector((state) => state.renderingState.cur_prompt); + const dispatch = useDispatch(); const [display_render_time, set_display_render_time] = useState(false); + const [display_generative_prompt, set_display_generative_prompt] = useState(false); + + const receive_generative_prompt = (e) => { + const msg = msgpack.decode(new Uint8Array(e.data)); + if (msg.path === '/model/is_generative') { + set_display_generative_prompt(msg.data === 'true'); + websocket.removeEventListener('message', receive_generative_prompt); + } + }; + websocket.addEventListener('message', receive_generative_prompt); const receive_temporal_dist = (e) => { const msg = msgpack.decode(new Uint8Array(e.data)); @@ -265,25 +277,44 @@ export function RenderControls() { ); }, }, + // Dreamfusion prompt editing + ...(display_generative_prompt + ? { + generative_prompt: { + label: 'Dreamfusion Prompt', + // value: "test string", + value: generative_prompt, + hint: 'Prompt used to guide Dreamfusion training', + onChange: (value) => { + dispatch_and_send( + websocket, + dispatch, + 'renderingState/new_prompt', + value, + ); + }, + }, + } + : {}), // Dynamic NeRF rendering time ...(display_render_time ? { - render_time: { - label: 'Render Timestep', - value: render_time, - min: 0, - max: 1, - step: 0.01, - onChange: (v) => { - dispatch_and_send( - websocket, - dispatch, - 'renderingState/render_time', - v, - ); - }, + render_time: { + label: 'Render Timestep', + value: render_time, + min: 0, + max: 1, + step: 0.01, + onChange: (v) => { + dispatch_and_send( + websocket, + dispatch, + 'renderingState/render_time', + v, + ); }, - } + }, + } : {}), }), [ @@ -294,6 +325,8 @@ export function RenderControls() { max_resolution, crop_enabled, target_train_util, + display_generative_prompt, + generative_prompt, render_time, display_render_time, websocket, // need to re-render when websocket changes to use the new websocket @@ -321,6 +354,7 @@ export function RenderControls() { crop_bg_color, crop_scale, crop_center, + display_generative_prompt, display_render_time, ]); diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 2c1541467..2613934d6 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -492,6 +492,21 @@ class ViewerState: has_temporal_distortion = getattr(graph, "temporal_distortion", None) is not None self.vis["model/has_temporal_distortion"].write(str(has_temporal_distortion).lower()) + prompt = getattr(graph, "prompt", None) + is_generative = prompt is not None + + if is_generative: + self.vis["model/is_generative"].write(str(is_generative).lower()) + + orig_prompt_written = self.vis["renderingState/cur_prompt"].read() is not None + if not orig_prompt_written: + self.vis["renderingState/cur_prompt"].write(str(prompt)) + self.vis["renderingState/new_prompt"].write(str(prompt)) + + new_prompt = self.vis["renderingState/new_prompt"].read() + if new_prompt != prompt: + setattr(graph, "prompt", str(new_prompt)) + is_training = self.vis["renderingState/isTraining"].read() self.step = step