diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index dc9aba07f..bba2f3df4 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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, ), ), diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index 4960fbe48..9c78513fc 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/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 diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index a3955712b..51e03cecb 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/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 diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index 187b1a60a..d7a94805e 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/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 diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index d3a658851..86e45e0e1 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index bbdd1099b..0abd1eaf3 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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