cleanup

nerfstudio/configs/method_configs.py

@@ -365,7 +365,7 @@ method_configs["dreamfusion"] = TrainerConfig(
             start_lambertian_training=1000,
             start_normals_training=500,
             opacity_loss_mult=0.001,
-            prompting_type="location_based",
+            positional_prompting="discrete",
             guidance_scale=100,
         ),
     ),

nerfstudio/fields/dreamfusion_field.py

@@ -29,7 +29,6 @@ from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.field_components.activations import trunc_exp
 from nerfstudio.field_components.field_heads import FieldHeadNames, PredNormalsFieldHead
 from nerfstudio.fields.base_field import Field
-from nerfstudio.utils import math
 
 try:
     import tinycudann as tcnn

nerfstudio/generative/stable_diffusion.py

@@ -199,32 +199,34 @@ class StableDiffusion(nn.Module):
         Returns:
             The loss
         """
-        image = F.interpolate(image, (IMG_DIM, IMG_DIM), mode="bilinear")
-        t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
-        latents = self.imgs_to_latent(image)
+        # Disable autocast when using diffusers pipeline.
+        with torch.autocast(device_type="cuda", enabled=False):
+            image = F.interpolate(image.half(), (IMG_DIM, IMG_DIM), mode="bilinear")
+            t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
+            latents = self.imgs_to_latent(image)
 
-        # predict the noise residual with unet, NO grad!
-        with torch.no_grad():
-            # add noise
-            noise = torch.randn_like(latents)
-            latents_noisy = self.scheduler.add_noise(latents, noise, t)  # type: ignore
-            # pred noise
-            latent_model_input = torch.cat([latents_noisy] * 2)
-            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)  # type: ignore
+                # pred noise
+                latent_model_input = torch.cat([latents_noisy] * 2)
+                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
 
-        # perform guidance
-        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-        noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond)
+            # perform guidance
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
-        # w(t), sigma_t^2
-        w = 1 - self.alphas[t]
+            # w(t), sigma_t^2
+            w = 1 - self.alphas[t]
 
-        grad = w * (noise_pred - noise)
-        grad = torch.nan_to_num(grad)
+            grad = w * (noise_pred - noise)
+            grad = torch.nan_to_num(grad)
 
-        if grad_scaler is not None:
-            latents = grad_scaler.scale(latents)
-        loss = _SDSGradient.apply(latents, grad)
+            if grad_scaler is not None:
+                latents = grad_scaler.scale(latents)
+            loss = _SDSGradient.apply(latents, grad)
 
         return loss
 

nerfstudio/generative/stable_diffusion_utils.py

@@ -16,47 +16,95 @@
 Stable diffusion utils
 """
 
+from torchtyping import TensorType
+from typing_extensions import Literal
 
-def get_text_embedding(batch, prompting_type, text_embeddings, sd_device):
-    """"""
+from nerfstudio.generative.stable_diffusion import StableDiffusion
 
-    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")
+class PositionalTextEmbeddings:
+    """Postional Prompts. Currently the following location based text embeddings are supported:
+        - discrete: Choose the embedding closest to the camera position
+        - interpolated: Interpolate between the embeddings based on the camera position
+        - off: Don't modify the text embedding based on the camera position
 
-    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
+    Args:
+        base_prompt: Prompt for base view
+        top_prompt: Prompt for top view
+        side_prompt: Prompt for side view
+        back_prompt: Prompt for back view
+        front_prompt: Prompt for front view
+        stable_diffusion: Instance of StableDiffusion
+        positional_prompting: how to incorporate position into prompt.
+    """
 
-    elif prompting_type == "interpolated":
-        horiz = batch["central"].to(sd_device)
-        vert = max(batch["vertical"].to(sd_device), 0)
+    def __init__(
+        self,
+        base_prompt: str,
+        top_prompt: str,
+        side_prompt: str,
+        back_prompt: str,
+        front_prompt: str,
+        stable_diffusion: StableDiffusion,
+        positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete",
+    ):
+        self.positional_prompting = positional_prompting
+        self.sd_device = stable_diffusion.device
+        self.sd = stable_diffusion
+        self.update_prompt(base_prompt, top_prompt, side_prompt, back_prompt, front_prompt)
 
-        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
+    def update_prompt(self, base_prompt: str, top_prompt: str, side_prompt: str, back_prompt: str, front_prompt: str):
+        """Update the text embeddings based on the new prompts.
 
-        text_embedding = text_embedding / 90.0
-        text_embedding = (vert * text_embedding + (90 - vert) * top) / 90.0
+        Args:
+            base_prompt: Prompt for base view
+            top_prompt: Prompt for top view
+            side_prompt: Prompt for side view
+            back_prompt: Prompt for back view
+            front_prompt: Prompt for front view
+        """
+        self.base_embed = self.sd.get_text_embeds(base_prompt, "")
+        self.top_embed = self.sd.get_text_embeds(top_prompt, "")
+        self.side_embed = self.sd.get_text_embeds(side_prompt, "")
+        self.back_embed = self.sd.get_text_embeds(back_prompt, "")
+        self.front_embed = self.sd.get_text_embeds(front_prompt, "")
 
-    elif prompting_type == "base":
-        text_embedding = base
+    def get_text_embedding(self, vertical_angle: TensorType[1], horizonal_angle: TensorType[1]):
+        """Get text embedding based on the position of the camera relative to the scene.
+        This trick is used in Dreamfusion (https://dreamfusion3d.github.io/).
 
-    return text_embedding
+        Args:
+            vertical_angle: vertical angle of the camera
+            horizonal_angle: horizonal angle of the camera
+        """
+
+        if self.positional_prompting == "discrete":
+            if vertical_angle < 40:
+                text_embedding = self.top_embed
+            elif 315 < horizonal_angle <= 45:
+                text_embedding = self.front_embed
+            elif 45 < horizonal_angle <= 135:
+                text_embedding = self.side_embed
+            elif 135 < horizonal_angle <= 225:
+                text_embedding = self.back_embed
+            else:  # horizonal_angle > 225 and horizonal_angle <= 315:
+                text_embedding = self.side_embed
+        elif self.positional_prompting == "interpolated":
+            horiz = horizonal_angle.to(self.sd_device)
+            vert = max(vertical_angle.to(self.sd_device), 0)
+
+            if 0 < horizonal_angle <= 90:
+                text_embedding = (horiz) * self.side_embed + (90 - horiz) * self.front_embed
+            elif 90 < horizonal_angle <= 180:
+                text_embedding = (horiz - 90) * self.back_embed + (180 - horiz) * self.side_embed
+            elif 180 < horizonal_angle <= 270:
+                text_embedding = (horiz - 180) * self.side_embed + (270 - horiz) * self.back_embed
+            else:  # 270 < horizonal_angle <= 360:
+                text_embedding = (horiz - 270) * self.front_embed + (360 - horiz) * self.side_embed
+
+            text_embedding = text_embedding / 90.0
+            text_embedding = (vert * text_embedding + (90 - vert) * self.top_embed) / 90.0
+        else:
+            text_embedding = self.base_embed
+
+        return text_embedding

nerfstudio/models/dreamfusion.py

@@ -23,8 +23,8 @@ from typing import Dict, List, Optional, Tuple, Type
 
 import numpy as np
 import torch
-from torch.cuda.amp.grad_scaler import GradScaler
 from torch.nn import Parameter
+from typing_extensions import Literal
 
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.engine.callbacks import (
@@ -36,7 +36,7 @@ from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.fields.density_fields import HashMLPDensityField
 from nerfstudio.fields.dreamfusion_field import DreamFusionField
 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.model_components.losses import (
     MSELoss,
     distortion_loss,
@@ -77,7 +77,7 @@ class DreamFusionModelConfig(ModelConfig):
     """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"""
@@ -135,8 +135,8 @@ class DreamFusionModelConfig(ModelConfig):
     """enables location based prompting"""
     interpolated_prompting: bool = False
     """enables interpolated location prompting"""
-    prompting_type: str = "location_based"
-    """choose between location_based, interpolated, base"""
+    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"
@@ -177,7 +177,7 @@ class DreamFusionModel(Model):
         self.target_transmittance = config.target_transmittance_start
         self.grad_scaler = kwargs["grad_scaler"]
 
-        self.prompting_type = config.prompting_type
+        self.positional_prompting = config.positional_prompting
         self.guidance_scale = config.guidance_scale
         self.top_prompt = config.top_prompt
         self.side_prompt = config.side_prompt
@@ -197,13 +197,15 @@ class DreamFusionModel(Model):
         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, ""),
-        }
+        self.text_embeddings = PositionalTextEmbeddings(
+            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,
+            stable_diffusion=self.sd,
+            positional_prompting=self.positional_prompting,
+        )
 
         # setting up fields
         self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res)
@@ -340,25 +342,26 @@ class DreamFusionModel(Model):
         param_groups["fields"] = list(self.field.parameters())
         return param_groups
 
-    def get_outputs(self, ray_bundle: RayBundle):
+    def get_outputs(self, ray_bundle: RayBundle):  # pylint: disable=too-many-statements
         # uniform sampling
         background_rgb = self.field.get_background_rgb(ray_bundle)
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         field_outputs = self.field(ray_samples, compute_normals=True)
-        weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
-        weights_list.append(weights)
-        ray_samples_list.append(ray_samples)
         density = field_outputs[FieldHeadNames.DENSITY]
 
-        if self.initialize_density and self.training:
+        if self.initialize_density:
             pos = ray_samples.frustums.get_positions()
             density_blob = (
                 self.config.init_density_strength
                 * self.density_strength
-                * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
+                * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None]
             )
             density += density_blob
 
+        weights = ray_samples.get_weights(density)
+        weights_list.append(weights)
+        ray_samples_list.append(ray_samples)
+
         accumulation = self.renderer_accumulation(weights)
         depth = self.renderer_depth(weights, ray_samples)
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
@@ -422,7 +425,7 @@ class DreamFusionModel(Model):
         else:
             outputs["train_output"] = outputs["rgb"]
 
-        if self.training or True:
+        if self.training:
             outputs["rendered_orientation_loss"] = orientation_loss(
                 weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
             )
@@ -464,15 +467,17 @@ class DreamFusionModel(Model):
 
         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,
+            )
 
-        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"]
@@ -480,13 +485,12 @@ class DreamFusionModel(Model):
             .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,
-            )
+        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)
 

nerfstudio/viewer/server/viewer_utils.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# pylint: disable=too-many-lines
+
 """Code to interface with the `vis/` (the JS viewer).
 """
 from __future__ import annotations