add prompt edditing to viewer

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)

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)

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,
   ]);
 

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