diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 92fe1966e..957650053 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -43,6 +43,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( ) from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind +from nerfstudio.models.dreamfusion import DreamFusionModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig @@ -62,6 +63,7 @@ descriptions = { "tensorf": "tensorf", "dnerf": "Dynamic-NeRF model. (slow)", "phototourism": "Uses the Phototourism data.", + "dreamfusion": "NeRF Model used in implementation of stable dreamfusion.", } method_configs["nerfacto"] = Config( @@ -249,6 +251,37 @@ method_configs["phototourism"] = Config( vis="viewer", ) +method_configs["dreamfusion"] = Config( + method_name="dreamfusion", + trainer=TrainerConfig( + steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True + ), + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + dataparser=NerfstudioDataParserConfig(), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=DreamFusionModelConfig(), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) + + AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[ # Don't show unparseable (fixed) arguments in helptext. tyro.conf.FlagConversionOff[ tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions) diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index 486b6011c..f9830eeed 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/nerfstudio/fields/dreamfusion_field.py @@ -70,9 +70,6 @@ class DreamFusionField(Field): geo_feat_dim: output geo feat dimensions num_layers_color: number of hidden layers for color network hidden_dim_color: dimension of hidden layers for color network - appearance_embedding_dim: dimension of appearance embedding - transient_embedding_dim: dimension of transient embedding - use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference spatial_distortion: spatial distortion to apply to the scene """ @@ -85,7 +82,6 @@ class DreamFusionField(Field): num_layers_color: int = 3, hidden_dim_color: int = 64, hidden_dim_transient: int = 64, - use_pred_normals: bool = False, spatial_distortion: Optional[SpatialDistortion] = None, ) -> None: super().__init__() @@ -94,7 +90,6 @@ class DreamFusionField(Field): self.geo_feat_dim = geo_feat_dim self.spatial_distortion = spatial_distortion - self.use_pred_normals = use_pred_normals num_levels = 16 max_res = 1024 @@ -103,14 +98,6 @@ class DreamFusionField(Field): features_per_level = 2 growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) - self.direction_encoding = tcnn.Encoding( - n_input_dims=3, - encoding_config={ - "otype": "SphericalHarmonics", - "degree": 4, - }, - ) - self.position_encoding = tcnn.Encoding( n_input_dims=3, encoding_config={"otype": "Frequency", "n_frequencies": 2}, @@ -137,19 +124,18 @@ class DreamFusionField(Field): ) # predicted normals - if self.use_pred_normals: - self.mlp_pred_normals = tcnn.Network( - n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims, - n_output_dims=hidden_dim_transient, - network_config={ - "otype": "FullyFusedMLP", - "activation": "ReLU", - "output_activation": "None", - "n_neurons": 64, - "n_hidden_layers": 2, - }, - ) - self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims) + self.mlp_pred_normals = tcnn.Network( + n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims, + n_output_dims=hidden_dim_transient, + network_config={ + "otype": "FullyFusedMLP", + "activation": "ReLU", + "output_activation": "None", + "n_neurons": 64, + "n_hidden_layers": 2, + }, + ) + self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims) self.mlp_head = tcnn.Network( n_input_dims=self.geo_feat_dim, @@ -190,22 +176,19 @@ class DreamFusionField(Field): outputs = {} directions = get_normalized_directions(ray_samples.frustums.directions) - directions_flat = directions.view(-1, 3) - d = self.direction_encoding(directions_flat) outputs_shape = ray_samples.frustums.directions.shape[:-1] # predicted normals - if self.use_pred_normals: - positions = ray_samples.frustums.get_positions() + positions = ray_samples.frustums.get_positions() - positions_flat = self.position_encoding(positions.view(-1, 3)) - pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + positions_flat = self.position_encoding(positions.view(-1, 3)) + pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) - x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions) - outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x) + x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions) + outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x) - h = [density_embedding.view(-1, self.geo_feat_dim)] + h = density_embedding.view(-1, self.geo_feat_dim) rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions) outputs.update({FieldHeadNames.RGB: rgb}) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py new file mode 100644 index 000000000..d2eb439f2 --- /dev/null +++ b/nerfstudio/generative/stable_diffusion.py @@ -0,0 +1,196 @@ + +# From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py + +import torch +from diffusers import AutoencoderKL, PNDMScheduler, UNet2DConditionModel +from rich.console import Console +from torch import nn +from transformers import CLIPTextModel, CLIPTokenizer, logging + +CONSOLE = Console(width=120) +logging.set_verbosity_error() +IMG_DIM = 512 + +class StableDiffusion(nn.Module): + """Stable Diffusion implementation + """ + + def __init__( + self, + device + ) -> None: + super().__init__() + + self.device = device + + try: + with open('./HF_TOKEN', 'r') as f: + self.token = f.read().replace('\n', '') + CONSOLE.print(f'Hugging face access token loaded.') + except FileNotFoundError as e: + raise AssertionError(f'Cannot read hugging face token, make sure to save your Hugging Face token at `./HF_TOKEN`') + + SD_SOURCE = "runwayml/stable-diffusion-v1-5" + CLIP_SOURCE = "openai/clip-vit-large-patch14" + + self.num_train_timesteps = 1000 + self.min_step = int(self.num_train_timesteps * 0.02) + self.max_step = int(self.num_train_timesteps * 0.98) + + CONSOLE.print(f'Loading Stable Diffusion model from {SD_SOURCE}.') + + self.auto_encoder = AutoencoderKL.from_pretrained(SD_SOURCE, subfolder="vae", use_auth_token=self.token).to(self.device) + self.unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", use_auth_token=self.token).to(self.device) + + CONSOLE.print(f'Loading CLIP model from {CLIP_SOURCE}.') + + self.tokenizer = CLIPTokenizer.from_pretrained(CLIP_SOURCE) + self.text_encoder = CLIPTextModel.from_pretrained(CLIP_SOURCE).to(self.device) + + self.scheduler = PNDMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=self.num_train_timesteps) + self.alphas = self.scheduler.alphas_cumprod.to(self.device) + + CONSOLE.print(f'Stable Diffusion loaded!') + + def get_text_embeds(self, prompt, negative_prompt): + # prompt, negative_prompt: [str] + + # Tokenize text and get embeddings + text_input = self.tokenizer(prompt, padding='max_length', max_length=self.tokenizer.model_max_length, truncation=True, return_tensors='pt') + + with torch.no_grad(): + text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0] + + # Do the same for unconditional embeddings + uncond_input = self.tokenizer(negative_prompt, padding='max_length', max_length=self.tokenizer.model_max_length, return_tensors='pt') + + with torch.no_grad(): + uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0] + + # Cat for final embeddings + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + return text_embeddings + + + def train_step(self, text_embeddings, pred_nerf, guidance_scale=100): + + # Upscale NeRF render for Stable Diffusion + pred_nerf = F.interpolate(pred_nerf, (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(pred_nerf) + + # 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) + # 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 (high scale from paper!) + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # w(t), sigma_t^2 + w = (1 - self.alphas[t]) + # w = self.alphas[t] ** 0.5 * (1 - self.alphas[t]) + grad = w * (noise_pred - noise) + + # clip grad for stable training? + # grad = grad.clamp(-10, 10) + grad = torch.nan_to_num(grad) + + # manually backward, since we omitted an item in grad and cannot simply autodiff. + # _t = time.time() + latents.backward(gradient=grad, retain_graph=True) + # torch.cuda.synchronize(); print(f'[TIME] guiding: backward {time.time() - _t:.4f}s') + + return 0 # dummy loss value + + def produce_latents(self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None): + + if latents is None: + latents = torch.randn((text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device) + + self.scheduler.set_timesteps(num_inference_steps) + + with torch.autocast('cuda'): + for t in self.scheduler.timesteps: + # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. + latent_model_input = torch.cat([latents] * 2) + + # predict the noise residual + with torch.no_grad(): + 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_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents)['prev_sample'] + + return latents + + def latents_to_img(self, latents): + + latents = 1 / 0.18215 * latents + + with torch.no_grad(): + imgs = self.auto_encoder.decode(latents).sample + + imgs = (imgs / 2 + 0.5).clamp(0, 1) + + return imgs + + def imgs_to_latent(self, imgs): + # imgs: [B, 3, H, W] + + imgs = 2 * imgs - 1 + + posterior = self.auto_encoder.encode(imgs).latent_dist + latents = posterior.sample() * 0.18215 + + return latents + + def prompt_to_img(self, prompts, negative_prompts='', num_inference_steps=50, guidance_scale=7.5, latents=None): + + prompts = [prompts] if isinstance(prompts, str) else prompts + negative_prompts = [negative_prompts] if isinstance(negative_prompts, str) else negative_prompts + text_embeddings = self.get_text_embeds(prompts, negative_prompts) + latents = self.produce_latents(text_embeddings, height=IMG_DIM, width=IMG_DIM, latents=latents, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale) # [1, 4, 64, 64] + + diffused_img = self.latents_to_img(latents) + diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy() + diffused_img = (diffused_img * 255).round().astype('uint8') + + return diffused_img + +def seed_everything(seed): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + +if __name__ == '__main__': + + import argparse + + import matplotlib.pyplot as plt + + parser = argparse.ArgumentParser() + parser.add_argument('prompt', type=str) + parser.add_argument('--negative', default='', type=str) + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--steps', type=int, default=50) + opt = parser.parse_args() + + seed_everything(opt.seed) + + device = torch.device('cuda') + + sd = StableDiffusion(device) + + imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.steps) + + plt.imsave('test_sd.png', imgs[0]) \ No newline at end of file diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 1946e4abc..a5529e5c8 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -182,7 +182,7 @@ def orientation_loss( """ w = weights n = normals - v = viewdirs + v = viewdirs * -1 n_dot_v = (n * v[..., None, :]).sum(axis=-1) return (w[..., 0] * torch.fmin(torch.zeros_like(n_dot_v), n_dot_v) ** 2).sum(dim=-1) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 4963e1dfc..55868cc04 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -30,19 +30,24 @@ from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.rays import RayBundle -from nerfstudio.field_components.encodings import NeRFEncoding, +# from nerfstudio.field_components.encodings import NeRFEncoding, from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.dreamfusion_field import DreamFusionField -from nerfstudio.model_components.losses import MSELoss +from nerfstudio.model_components.losses import ( + MSELoss, + orientation_loss, + pred_normal_loss, +) from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler from nerfstudio.model_components.renderers import ( AccumulationRenderer, DepthRenderer, + NormalsRenderer, RGBRenderer, ) from nerfstudio.model_components.scene_colliders import AABBBoxCollider from nerfstudio.models.base_model import Model, ModelConfig -from nerfstudio.utils import colormaps, colors, misc +from nerfstudio.utils import colormaps, colors, math, misc @dataclass @@ -54,6 +59,10 @@ class DreamFusionModelConfig(ModelConfig): num_samples: int = 256 """Number of samples in field evaluation""" + orientation_loss_mult: float = 0.0001 + """Orientation loss multipier on computed normals.""" + pred_normal_loss_mult: float = 0.001 + """Predicted normal loss multiplier.""" class DreamFusionModel(Model): @@ -63,6 +72,8 @@ class DreamFusionModel(Model): config: DreamFusion configuration to instantiate model """ + config: DreamFusionModelConfig + def __init__( self, config: DreamFusionModelConfig, @@ -87,6 +98,7 @@ class DreamFusionModel(Model): self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer() + self.renderer_normals = NormalsRenderer() # losses self.rgb_loss = MSELoss() @@ -108,14 +120,46 @@ class DreamFusionModel(Model): def get_outputs(self, ray_bundle: RayBundle): # uniform sampling ray_samples_uniform = self.sampler_uniform(ray_bundle) - field_outputs = self.field(ray_samples_uniform) + field_outputs = self.field(ray_samples_uniform, compute_normals=True) weights = ray_samples_uniform.get_weights(field_outputs[FieldHeadNames.DENSITY]) accumulation = self.renderer_accumulation(weights) depth = self.renderer_depth(weights, ray_samples_uniform) rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + + # lambertian shading + light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) + light_d = math.safe_normalize(light_d) + + ratio = 0.1 + lambertian = ratio + (1 - ratio) * (pred_normals @ light_d).clamp(min=0) # [N,] + + shaded = lambertian.unsqueeze(-1).repeat(1, 3) + shaded_albedo = rgb * lambertian.unsqueeze(-1) + + outputs = { + "rgb": rgb, + "accumulation": accumulation, + "depth": depth, + "normals": normals, + "pred_normals": pred_normals, + "shaded": shaded, + "shaded_albedo": shaded_albedo, + } + + if self.training: + outputs["rendered_orientation_loss"] = orientation_loss( + weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions + ) + + outputs["rendered_pred_normal_loss"] = pred_normal_loss( + weights.detach(), + field_outputs[FieldHeadNames.NORMALS].detach(), + field_outputs[FieldHeadNames.PRED_NORMALS], + ) - outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth} return outputs def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: @@ -128,6 +172,16 @@ class DreamFusionModel(Model): loss_dict = {"rgb_loss": rgb_loss} loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) + if self.training: + # orientation loss for computed normals + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( + outputs["rendered_orientation_loss"] + ) + + # ground truth supervision for normals + loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean( + outputs["rendered_pred_normal_loss"] + ) return loss_dict def get_image_metrics_and_images( diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index b1e772b03..15d139054 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -99,7 +99,7 @@ class NerfactoModelConfig(ModelConfig): distortion_loss_mult: float = 0.002 """Distortion loss multiplier.""" orientation_loss_mult: float = 0.0001 - """Orientation loss multipier on computed noramls.""" + """Orientation loss multipier on computed normals.""" pred_normal_loss_mult: float = 0.001 """Predicted normal loss multiplier.""" use_proposal_weight_anneal: bool = True diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index d5ed57f94..f5a0f8219 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -191,3 +191,8 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor: """ return torch.exp(-0.5 * x_vars) * torch.sin(x_means) + + +# From https://github.com/ashawkey/stable-dreamfusion +def safe_normalize(x, eps=1e-20): + return x / torch.sqrt(torch.clamp(torch.sum(x * x, -1, keepdim=True), min=eps)) diff --git a/pyproject.toml b/pyproject.toml index 2d22d0684..ebd472649 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,8 @@ dependencies = [ "rich==12.5.1", "tensorboard==2.9.0", "torch==1.12.1", + "diffusers==0.9.0", + "transformers==4.24.0", "torchmetrics[image]>=0.9.3", "torchtyping>=0.1.4", "torchvision>=0.13.0",