Sigmoid on background, prompt in config

nerfstudio/fields/dreamfusion_field.py

@@ -143,7 +143,7 @@ class DreamFusionField(Field):
             network_config={
                 "otype": "FullyFusedMLP",
                 "activation": "ReLU",
-                "output_activation": "None",
+                "output_activation": "Sigmoid",
                 "n_neurons": 64,
                 "n_hidden_layers": 2,
             },
@@ -193,7 +193,6 @@ class DreamFusionField(Field):
         background_rgb = self.mlp_background_color(positions_flat).view(*outputs_shape, -1).to(directions)
         return background_rgb
 
-
     def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
         assert density_embedding is not None
         outputs = {}

nerfstudio/generative/stable_diffusion.py

@@ -1,4 +1,3 @@
-
 # From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py
 
 import matplotlib.pyplot as plt
@@ -15,24 +14,23 @@ logging.set_verbosity_error()
 IMG_DIM = 512
 const_scale = 0.18215
 
-class StableDiffusion(nn.Module):
-    """Stable Diffusion implementation
-    """
 
-    def __init__(
-        self,
-        device
-    ) -> None:
+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.')
+            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`')
+            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"
@@ -41,32 +39,49 @@ class StableDiffusion(nn.Module):
         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}.')
+        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)
+        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}.')
+        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.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!')
+        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')
+        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')
+        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]
@@ -79,10 +94,10 @@ class StableDiffusion(nn.Module):
 
         nerf_output_np = nerf_output.view(3, 64, 64).permute(1, 2, 0).detach().cpu().numpy()
         nerf_output_np = np.clip(nerf_output_np, 0.0, 1.0)
-        plt.imsave('nerf_output.png', nerf_output_np)
-        nerf_output = F.interpolate(nerf_output, (IMG_DIM, IMG_DIM), mode='bilinear')
+        plt.imsave("nerf_output.png", nerf_output_np)
+        nerf_output = F.interpolate(nerf_output, (IMG_DIM, IMG_DIM), mode="bilinear")
 
-        print(torch.max(nerf_output), torch.min(nerf_output))
+        # print(torch.max(nerf_output), torch.min(nerf_output))
 
         t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
         latents = self.imgs_to_latent(nerf_output)
@@ -100,7 +115,7 @@ class StableDiffusion(nn.Module):
         noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
         # w(t), sigma_t^2
-        w = (1 - self.alphas[t])
+        w = 1 - self.alphas[t]
         # w = self.alphas[t] ** 0.5 * (1 - self.alphas[t])
         grad = w * (noise_pred - noise)
 
@@ -120,29 +135,33 @@ class StableDiffusion(nn.Module):
 
         return noise_loss, latents, grad
 
-    def produce_latents(self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None):
+    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)
+            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'):
+        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']
+                    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']
-        
+                latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"]
+
         return latents
 
     def latents_to_img(self, latents):
@@ -153,7 +172,7 @@ class StableDiffusion(nn.Module):
             imgs = self.auto_encoder.decode(latents).sample
 
         imgs = (imgs / 2 + 0.5).clamp(0, 1)
-        
+
         return imgs
 
     def imgs_to_latent(self, imgs):
@@ -165,40 +184,49 @@ class StableDiffusion(nn.Module):
 
         return latents
 
-    def prompt_to_img(self, prompts, negative_prompts='', num_inference_steps=50, guidance_scale=7.5, latents=None):
-       
+    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]
-        
+        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')
+        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__':
+
+if __name__ == "__main__":
 
     import argparse
 
     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)
+    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')
+    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])
+    plt.imsave("test_sd.png", imgs[0])

nerfstudio/models/dreamfusion.py

@@ -72,6 +72,8 @@ class DreamFusionModelConfig(ModelConfig):
     """Randomizes light source per output."""
     initialize_density: bool = True
     """Initialize density in center of scene."""
+    sphere_collider: bool = True
+    """Use spherical collider instead of box"""
 
 
 class DreamFusionModel(Model):
@@ -151,9 +153,10 @@ class DreamFusionModel(Model):
         self.lpips = LearnedPerceptualImagePatchSimilarity()
 
         # colliders
-        if self.config.enable_collider:
+        if self.config.sphere_collider:
+            self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 0.8)
+        else:
             self.collider = AABBBoxCollider(scene_box=self.scene_box)
-        self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 0.8)
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
         param_groups = {}
@@ -178,7 +181,7 @@ class DreamFusionModel(Model):
         accumulation = self.renderer_accumulation(weights)
         depth = self.renderer_depth(weights, ray_samples_uniform)
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
-        rgb = torch.clamp(rgb, 0.0, 1.0)
+        # rgb = torch.clamp(rgb, 0.0, 1.0)
 
         accum_mask = torch.clamp((torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0)
         accum_mask_inv = 1.0 - accum_mask
@@ -194,9 +197,7 @@ class DreamFusionModel(Model):
         }
 
         normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
-        pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)#.detach()
-
-
+        pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)  # .detach()
 
         # lambertian shading
         if self.config.random_light_source:

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -23,9 +23,11 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     """target class to instantiate"""
     datamanager: DreamFusionDataManagerConfig = DreamFusionDataManagerConfig()
     """specifies the datamanager config"""
+    prompt: str = "A high quality photo of a pineapple"
+    """prompt for stable dreamfusion"""
+
 
 class DreamfusionPipeline(VanillaPipeline):
-    
     def __init__(
         self,
         config: DreamfusionPipelineConfig,
@@ -34,15 +36,9 @@ class DreamfusionPipeline(VanillaPipeline):
         world_size: int = 1,
         local_rank: int = 0,
     ):
-        super().__init__(
-            config,
-            device,
-            test_mode,
-            world_size,
-            local_rank
-        )
+        super().__init__(config, device, test_mode, world_size, local_rank)
         self.sd = StableDiffusion(device)
-        self.text_embedding = self.sd.get_text_embeds("A high quality photo of a squirrel.", "")
+        self.text_embedding = self.sd.get_text_embeds(config.prompt, "")
 
     @profiler.time_function
     def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers):
@@ -64,15 +60,15 @@ class DreamfusionPipeline(VanillaPipeline):
 
         accumulation = model_outputs["accumulation"].view(64, 64).detach().cpu().numpy()
         accumulation = np.clip(accumulation, 0.0, 1.0)
-        plt.imsave('nerf_accumulation.jpg', accumulation)
+        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)
+        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)
+        plt.imsave("nerf_textureless.jpg", shaded)
 
         sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output)
 
@@ -92,4 +88,4 @@ class DreamfusionPipeline(VanillaPipeline):
         grad_scaler.update()
         optimizers.scheduler_step_all(step)
 
-        return model_outputs, loss_dict, metrics_dict
+        return model_outputs, loss_dict, metrics_dict