Dreamfusion pipeline

nerfstudio/configs/method_configs.py

@@ -54,6 +54,7 @@ from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig
 from nerfstudio.models.tensorf import TensoRFModelConfig
 from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig
 from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
+from nerfstudio.pipelines.dreamfusion_pipeline import DreamfusionPipelineConfig
 from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig
 
 method_configs: Dict[str, Config] = {}
@@ -261,7 +262,7 @@ method_configs["dreamfusion"] = Config(
     trainer=TrainerConfig(
         steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
     ),
-    pipeline=VanillaPipelineConfig(
+    pipeline=DreamfusionPipelineConfig(
         datamanager=DreamFusionDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -237,13 +237,12 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
         self.eval_cameras, _, _ = random_train_pose(
             self.config.num_eval_angles,
             self.config.eval_resolution,
-            device=device,
-            # radius_mean=1,
-            # radius_std=0.1,
-            # central_rotation_range=[180, 180],
-            # vertical_rotation_range=[0, 0],
-            # focal_range=[1, 1],
-            # jitter_std=0.1,
+            device=self.device,
+            radius_mean=1.5,
+            radius_std=0,
+            focal_range=[1, 1],
+            vertical_rotation_range=[0, 20],
+            # jitter_std=0,
         )
 
         self.train_dataset = TrivialDataset(self.eval_cameras)
@@ -256,14 +255,16 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
 
         self.train_count += 1
 
-        if step > 2000:
-            cameras, _, _ = random_train_pose(
-                self.config.train_images_per_batch, self.config.train_resolution, device=self.device
-            )
+        # TODO Reimplement when cameras are fully working
+        # if step > 2000:
+        #     cameras, _, _ = random_train_pose(
+        #         self.config.train_images_per_batch, self.config.train_resolution, device=self.device
+        #     )
 
-            ray_bundle = cameras.generate_rays(torch.tensor([[i] for i in range(self.config.train_images_per_batch)]))
-            return ray_bundle, {"initialization": False}
+        #     ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten()
+        #     return ray_bundle, {"initialization": False}
 
+        # TODO below
         cameras, vertical_rotation, central_rotation = random_train_pose(
             self.config.train_images_per_batch,
             self.config.train_resolution,
@@ -271,10 +272,13 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
             radius_mean=1.5,
             radius_std=0,
             focal_range=[1, 1],
-            vertical_rotation_range=[-180, 180],
-            jitter_std=0,
+            # vertical_rotation_range=[-180, 180],
+            # jitter_std=0,
         )
-        ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten()
+        # ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten()
+
+        camera_idx = torch.randint(0, self.eval_cameras.shape[0], [1], dtype=torch.long, device=self.device)
+        ray_bundle = self.eval_cameras.generate_rays(camera_idx).flatten()
 
         return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation, "initialization": True}
 

nerfstudio/engine/trainer.py

@@ -311,13 +311,18 @@ class Trainer:
         """
         self.optimizers.zero_grad_all()
         cpu_or_cuda_str = self.device.split(":")[0]
-        with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision):
-            _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step)
+
+        if hasattr(self.pipeline, 'custom_step'):
+            _, loss_dict, metrics_dict = self.pipeline.custom_step(step=step, grad_scaler=self.grad_scaler, optimizers=self.optimizers)
             loss = functools.reduce(torch.add, loss_dict.values())
-        # self.grad_scaler.scale(loss).backward()  # type: ignore
-        # self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
-        # self.grad_scaler.update()
-        self.optimizers.scheduler_step_all(step)
+        else:
+            with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision):
+                _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step)
+                loss = functools.reduce(torch.add, loss_dict.values())
+            self.grad_scaler.scale(loss).backward()  # type: ignore
+            self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
+            self.grad_scaler.update()
+            self.optimizers.scheduler_step_all(step)
 
         # Merging loss and metrics dict into a single output.
         return loss, loss_dict, metrics_dict

nerfstudio/generative/stable_diffusion.py

@@ -2,6 +2,7 @@
 # From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py
 
 import matplotlib.pyplot as plt
+import numpy as np
 import torch
 import torch.nn.functional as F
 from diffusers import AutoencoderKL, PNDMScheduler, UNet2DConditionModel
@@ -75,7 +76,9 @@ class StableDiffusion(nn.Module):
 
     def sds_loss(self, text_embeddings, nerf_output, guidance_scale=100):
 
-        plt.imsave('nerf_output.png', nerf_output.view(3, 64, 64).permute(1, 2, 0).detach().cpu().numpy())
+        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')
 
         t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
@@ -102,13 +105,17 @@ class StableDiffusion(nn.Module):
         # grad = grad.clamp(-10, 10)
         grad = torch.nan_to_num(grad)
 
+        # noise_difference = self.latents_to_img(grad)
+        # plt.imsave('noise_difference.png', noise_difference.view(3, 512, 512).permute(1, 2, 0).detach().cpu().numpy())
+
         # manually backward, since we omitted an item in grad and cannot simply autodiff.
         # _t = time.time()
-        latents.backward(gradient=grad, retain_graph=True)
 
-        rgb_loss = torch.square(noise_pred - noise)
+        noise_loss = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise)))
 
-        return rgb_loss
+        # print('SDS LOSS:', noise_loss)
+
+        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):
 

nerfstudio/models/dreamfusion.py

@@ -38,7 +38,6 @@ from nerfstudio.engine.callbacks import (
 # from nerfstudio.field_components.encodings import NeRFEncoding,
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.fields.dreamfusion_field import DreamFusionField
-from nerfstudio.generative.stable_diffusion import StableDiffusion
 from nerfstudio.model_components.losses import (
     MSELoss,
     orientation_loss,
@@ -73,8 +72,6 @@ class DreamFusionModelConfig(ModelConfig):
     """Randomizes light source per output."""
     initialize_density: bool = True
     """Initialize density in center of scene."""
-    prompt: str = "A high quality photo of a pineaple."
-    """Text prompt to optimize scene for."""
 
 
 class DreamFusionModel(Model):
@@ -95,8 +92,6 @@ class DreamFusionModel(Model):
         self.initialize_density = config.initialize_density
         self.train_normals = False
         super().__init__(config=config, **kwargs)
-        self.sd = StableDiffusion(kwargs['device'])
-        self.text_embedding = self.sd.get_text_embeds(config.prompt, "")
 
     def get_training_callbacks(
         self, training_callback_attributes: TrainingCallbackAttributes
@@ -172,7 +167,7 @@ class DreamFusionModel(Model):
         density = field_outputs[FieldHeadNames.DENSITY]
         if self.initialize_density:
             pos = ray_samples_uniform.frustums.get_positions()
-            density_blob = 10 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
+            density_blob = 50 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
             density += density_blob
 
         weights = ray_samples_uniform.get_weights(density)
@@ -224,16 +219,9 @@ class DreamFusionModel(Model):
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
-        # TODO: modify to not take in data images, need to get diffusion loss here
         # Scaling metrics by coefficients to create the losses.
-        device = outputs["rgb"].device
 
-        # change to key 'shaded_albedo'
-        albedo_output = outputs["rgb"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
-
-        rgb_loss = self.sd.sds_loss(self.text_embedding, albedo_output)
-
-        loss_dict = {"rgb_loss": rgb_loss}
+        loss_dict = {}
         loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)
         if self.train_normals:
             # orientation loss for computed normals

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass, field
+from typing import Type
+
+from torch.cuda.amp.grad_scaler import GradScaler
+from typing_extensions import Literal
+
+from nerfstudio.data.datamanagers.dreamfusion_datamanager import (
+    DreamFusionDataManagerConfig,
+)
+from nerfstudio.engine.optimizers import Optimizers
+from nerfstudio.generative.stable_diffusion import StableDiffusion
+from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig
+from nerfstudio.utils import profiler
+
+
+@dataclass
+class DreamfusionPipelineConfig(VanillaPipelineConfig):
+    """Configuration for pipeline instantiation"""
+
+    _target: Type = field(default_factory=lambda: DreamfusionPipeline)
+    """target class to instantiate"""
+    datamanager: DreamFusionDataManagerConfig = DreamFusionDataManagerConfig()
+    """specifies the datamanager config"""
+
+class DreamfusionPipeline(VanillaPipeline):
+    
+    def __init__(
+        self,
+        config: DreamfusionPipelineConfig,
+        device: str,
+        test_mode: Literal["test", "val", "inference"] = "val",
+        world_size: int = 1,
+        local_rank: int = 0,
+    ):
+        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 pineapple.", "")
+
+    @profiler.time_function
+    def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers):
+        """This function gets your training loss dict. This will be responsible for
+        getting the next batch of data from the DataManager and interfacing with the
+        Model class, feeding the data to the model's forward function.
+
+        Args:
+            step: current iteration step to update sampler if using DDP (distributed)
+        """
+        if self.world_size > 1 and step:
+            assert self.datamanager.train_sampler is not None
+            self.datamanager.train_sampler.set_epoch(step)
+        ray_bundle, batch = self.datamanager.next_train(step)
+        model_outputs = self.model(ray_bundle)
+
+        # Just uses albedo for now
+        albedo_output = model_outputs["rgb"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
+
+        sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output)
+
+        grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True)
+        optimizers.optimizer_scaler_step_all(grad_scaler)
+        grad_scaler.update()
+        optimizers.scheduler_step_all(step)
+
+        metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
+        loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
+        loss_dict["sds_loss"] = sds_loss
+
+        return model_outputs, loss_dict, metrics_dict