From 0080e60ba5044cb16660e4514d81c3bfbc32f54b Mon Sep 17 00:00:00 2001 From: "Abhik Ahuja ahujaabhik@berkeley.edu" Date: Thu, 15 Dec 2022 06:15:24 +0000 Subject: [PATCH] Tuning, more config changes --- nerfstudio/generative/stable_diffusion.py | 2 +- nerfstudio/models/dreamfusion.py | 14 ++++++--- nerfstudio/pipelines/dreamfusion_pipeline.py | 32 ++++++++++++-------- 3 files changed, 31 insertions(+), 17 deletions(-) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 22e0652ec..f5e83d7ca 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -117,7 +117,7 @@ class StableDiffusion(nn.Module): # diffused_img = self.latents_to_img(latents_noisy - noise_pred) # diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy().reshape((512, 512, 3)) # diffused_img = (diffused_img * 255).round().astype("uint8") - # plt.imsave("sd_img.png", diffused_img) + # plt.imsave("sdimg.png", diffused_img) # w(t), sigma_t^2 w = 1 - self.alphas[t] diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 3bfe1372d..02046f6bc 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -66,12 +66,14 @@ class DreamFusionModelConfig(ModelConfig): """Number of samples in field evaluation""" orientation_loss_mult: float = 0.01 """Orientation loss multipier on computed normals.""" - pred_normal_loss_mult: float = 0.0003 + pred_normal_loss_mult: float = 0.001 """Predicted normal loss multiplier.""" random_light_source: bool = True """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" + init_density_strength: float = 50.0 + """Initial strength of center density""" sphere_collider: bool = True """Use spherical collider instead of box""" @@ -184,9 +186,13 @@ class DreamFusionModel(Model): field_outputs = self.field(ray_samples_uniform, compute_normals=True) density = field_outputs[FieldHeadNames.DENSITY] - if self.initialize_density: + if self.initialize_density and self.training: pos = ray_samples_uniform.frustums.get_positions() - density_blob = 50 * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] + density_blob = ( + self.config.init_density_strength + * self.density_strength + * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] + ) density += density_blob weights = ray_samples_uniform.get_weights(density) @@ -213,7 +219,7 @@ class DreamFusionModel(Model): pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) # .detach() # lambertian shading - if self.config.random_light_source: + if self.config.random_light_source and self.training: # and if training light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) else: light_d = ray_bundle.origins[0] diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 1a5845fb9..950eedfbc 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -25,12 +25,9 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """specifies the datamanager config""" prompt: str = "A high quality photo of a pineapple" """prompt for stable dreamfusion""" + location_based_prompting: bool = True """enables location based prompting""" - opacity_penalty: bool = True - """enables penalty to encourage transparent scenes""" - target_transmittance: float = 0.25 - """target transmittance for opacity penalty""" side_prompt: str = ", side view" """appended to prompt for side view""" front_prompt: str = ", front view" @@ -38,6 +35,14 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): back_prompt: str = ", back view" """appended to prompt for back view""" + opacity_penalty: bool = True + """enables penalty to encourage transparent scenes""" + target_transmittance: float = 0.5 + """target transmittance for opacity penalty""" + + guidance_scale: float = 100 + """guidance scale for sds loss""" + class DreamfusionPipeline(VanillaPipeline): def __init__( @@ -51,11 +56,12 @@ class DreamfusionPipeline(VanillaPipeline): super().__init__(config, device, test_mode, world_size, local_rank) self.generative = True self.sd = StableDiffusion(device) - self.base_text_embedding = self.sd.get_text_embeds(config.prompt, "") if config.location_based_prompting: 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}", "") + else: + self.base_text_embedding = self.sd.get_text_embeds(config.prompt, "") @profiler.time_function def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers): @@ -78,13 +84,13 @@ class DreamfusionPipeline(VanillaPipeline): accumulation = np.clip(accumulation, 0.0, 1.0) plt.imsave("nerf_accumulation.jpg", accumulation) - background = model_outputs["background"].view(64, 64, 3).detach().cpu().numpy() - background = np.clip(background, 0.0, 1.0) - plt.imsave("nerf_background.jpg", background) + # background = model_outputs["background"].view(64, 64, 3).detach().cpu().numpy() + # background = np.clip(background, 0.0, 1.0) + # plt.imsave("nerf_background.jpg", background) - shaded = model_outputs["shaded"].view(64, 64, 3).detach().cpu().numpy() - shaded = np.clip(shaded, 0.0, 1.0) - plt.imsave("nerf_textureless.jpg", shaded) + # shaded = model_outputs["shaded"].view(64, 64, 3).detach().cpu().numpy() + # shaded = np.clip(shaded, 0.0, 1.0) + # plt.imsave("nerf_textureless.jpg", shaded) if self.config.location_based_prompting: if batch["central"] > 315 or batch["central"] <= 45: @@ -98,7 +104,9 @@ class DreamfusionPipeline(VanillaPipeline): else: text_embedding = self.base_text_embedding - sds_loss, latents, grad = self.sd.sds_loss(text_embedding, train_output) + sds_loss, latents, grad = self.sd.sds_loss( + text_embedding, train_output, guidance_scale=self.config.guidance_scale + ) # TODO: opacity penalty using transmittance, not accumultation if self.config.opacity_penalty: accum_mean = np.mean(1.0 - accumulation)