merging

nerfstudio/configs/method_configs.py

@@ -47,6 +47,7 @@ from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataPars
 from nerfstudio.data.dataparsers.phototourism_dataparser import (
     PhototourismDataParserConfig,
 )
+from nerfstudio.engine.dreamfusion_trainer import DreamfusionTrainerConfig
 from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
 from nerfstudio.engine.schedulers import SchedulerConfig
 from nerfstudio.engine.trainer import TrainerConfig
@@ -332,11 +333,12 @@ method_configs["phototourism"] = TrainerConfig(
     vis="viewer",
 )
 
-method_configs["dreamfusion"] = ExperimentConfig(
+method_configs["dreamfusion"] = DreamfusionTrainerConfig(
     method_name="dreamfusion",
-    trainer=TrainerConfig(
-        steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
-    ),
+    steps_per_eval_batch=500,
+    steps_per_save=2000,
+    max_num_iterations=30000,
+    mixed_precision=True,
     pipeline=DreamfusionPipelineConfig(
         datamanager=DreamFusionDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -43,6 +43,7 @@ CONSOLE = Console(width=120)
 class TrivialDataset(InputDataset):
     """A trivial dataset with blank images for the viewer"""
 
+    # pylint: disable=super-init-not-called
     def __init__(self, cameras: Cameras):
         self.size = cameras.size
         self.cameras = cameras
@@ -159,6 +160,16 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig):
     """Number of images per batch for training"""
     eval_images_per_batch: int = 1
     """Number of images per batch for evaluation"""
+    radius_mean: float = 1.6
+    """Mean radius of camera orbit"""
+    radius_std: float = 0.2
+    """Std of radius of camera orbit"""
+    focal_range: Tuple[float, float] = (0.75, 1.35)
+    """Range of focal length"""
+    vertical_rotation_range: Tuple[float, float] = (-90, 30)
+    """Range of vertical rotation"""
+    jitter_std: float = 0.05
+    """Std of camera direction jitter, so we don't just point the cameras towards the center every time"""
 
 
 class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-method
@@ -176,6 +187,7 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
 
     config: DreamFusionDataManagerConfig
 
+    # pylint: disable=super-init-not-called
     def __init__(
         self,
         config: DreamFusionDataManagerConfig,
@@ -194,20 +206,21 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
         self.test_split = "test" if test_mode in ["test", "inference"] else "val"
         self.dataparser = self.config.dataparser.setup()
 
-        self.eval_cameras, _, _ = random_train_pose(
+        cameras, _, _ = random_train_pose(
             self.config.num_eval_angles,
             self.config.eval_resolution,
             device=self.device,
-            radius_mean=1.5,
-            radius_std=0,
-            focal_range=[1, 1],
-            vertical_rotation_range=[-30, 30],
-            # jitter_std=0,
+            radius_mean=self.config.radius_mean,
+            radius_std=self.config.radius_std,
+            focal_range=self.config.focal_range,
+            vertical_rotation_range=self.config.vertical_rotation_range,
+            jitter_std=self.config.jitter_std,
         )
 
-        self.train_dataset = TrivialDataset(self.eval_cameras)
-        self.eval_dataset = TrivialDataset(self.eval_cameras)
+        self.train_dataset = TrivialDataset(cameras)
+        self.eval_dataset = TrivialDataset(cameras)
 
+        # pylint: disable=non-parent-init-called
         DataManager.__init__(self)
 
     def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
@@ -215,13 +228,15 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
 
         self.train_count += 1
 
-        # TODO Reimplement when cameras are fully working
+        # # 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(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()
         #     return ray_bundle, {"initialization": False}
 
         # TODO below
@@ -229,16 +244,16 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
             self.config.train_images_per_batch,
             self.config.train_resolution,
             device=self.device,
-            radius_mean=1.5,
-            radius_std=0,
-            focal_range=[1, 1],
-            # vertical_rotation_range=[-180, 180],
-            # jitter_std=0,
+            radius_mean=self.config.radius_mean,
+            radius_std=self.config.radius_std,
+            focal_range=self.config.focal_range,
+            vertical_rotation_range=self.config.vertical_rotation_range,
+            jitter_std=self.config.jitter_std,
         )
-        # 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()
+        # 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}
 
@@ -258,11 +273,11 @@ class DreamFusionDataManager(VanillaDataManager):  # pylint: disable=abstract-me
             self.config.eval_images_per_batch,
             self.config.eval_resolution,
             device=self.device,
-            radius_mean=1.5,
-            radius_std=0,
-            focal_range=[1, 1],
-            vertical_rotation_range=[-180, 180],
-            jitter_std=0,
+            radius_mean=self.config.radius_mean,
+            radius_std=self.config.radius_std,
+            focal_range=self.config.focal_range,
+            vertical_rotation_range=self.config.vertical_rotation_range,
+            jitter_std=self.config.jitter_std,
         )
         ray_bundle = cameras.generate_rays(torch.tensor([[i] for i in range(self.config.train_images_per_batch)]))
 

nerfstudio/engine/dreamfusion_trainer.py

@@ -0,0 +1,86 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dreamfusion trainer"""
+
+from dataclasses import dataclass, field
+from typing import Dict, Tuple, Type
+
+import torch
+
+from nerfstudio.engine.trainer import Trainer, TrainerConfig
+from nerfstudio.pipelines.dreamfusion_pipeline import (
+    DreamfusionPipeline,
+    DreamfusionPipelineConfig,
+)
+from nerfstudio.utils import profiler
+
+
+@dataclass
+class DreamfusionTrainerConfig(TrainerConfig):
+    """Configuration for trainer instantiation"""
+
+    _target: Type = field(default_factory=lambda: DreamfusionTrainer)
+    """target class to instantiate"""
+    pipeline: DreamfusionPipelineConfig = DreamfusionPipelineConfig()
+    """specifies the pipeline config"""
+    gradient_accumulation_steps: int = 4
+    """number of gradient accumulation steps before optimizer step"""
+
+
+class DreamfusionTrainer(Trainer):
+    """Dreamfusion trainer"""
+
+    pipeline: DreamfusionPipeline
+    config: DreamfusionTrainerConfig
+
+    def __init__(self, config: DreamfusionTrainerConfig, local_rank: int = 0, world_size: int = 1):
+        assert isinstance(config, DreamfusionTrainerConfig)
+        Trainer.__init__(self, config=config, local_rank=local_rank, world_size=world_size)
+
+    @profiler.time_function
+    def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """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)
+        """
+        for i in range(self.config.gradient_accumulation_steps):
+            self.optimizers.zero_grad_all()
+            model_outputs, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step)
+            latents = model_outputs["latents"]
+            grad = model_outputs["grad"]
+            self.grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True)
+            normals_loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"]
+            if self.pipeline.config.alphas_penalty:
+                normals_loss += loss_dict["alphas_loss"] * self.config.pipeline.alphas_scale
+            if self.pipeline.config.opacity_penalty:
+                normals_loss += loss_dict["opacity_loss"] * self.config.pipeline.opacity_scale
+            self.grad_scaler.scale(normals_loss).backward()  # type: ignore
+            # I noticed that without these del statements and the empty cache, I was getting OOM errors on the device
+            # running stable diffusion, since these variables aren't deleted (normally they would be freed when we
+            # exit the frame since they no longer are being referenced, but now we loop back instead of exiting the
+            # frame)
+            if i != self.config.gradient_accumulation_steps - 1:
+                del normals_loss, latents, grad, model_outputs, loss_dict, metrics_dict
+            torch.cuda.empty_cache()
+
+        self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
+
+        self.grad_scaler.update()
+        self.optimizers.scheduler_step_all(step)
+
+        return torch.tensor(0.0), loss_dict, metrics_dict

nerfstudio/engine/trainer.py

@@ -360,20 +360,13 @@ class Trainer:
         """
         self.optimizers.zero_grad_all()
         cpu_or_cuda_str = self.device.split(":")[0]
-
-        if hasattr(self.pipeline, "custom_step"):
-            _, loss_dict, metrics_dict = self.pipeline.custom_step(
-                step=step, grad_scaler=self.grad_scaler, optimizers=self.optimizers
-            )
+        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())
-        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)
+        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/fields/dreamfusion_field.py

@@ -17,31 +17,18 @@ Field for dreamfusion model
 """
 
 
-from typing import Dict, Optional, Tuple
+from typing import Optional
 
 import numpy as np
 import torch
-from torch import nn
 from torch.nn.parameter import Parameter
 from torchtyping import TensorType
 
 from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.field_components.activations import trunc_exp
-from nerfstudio.field_components.embedding import Embedding
-from nerfstudio.field_components.encodings import Encoding, HashEncoding, SHEncoding
-from nerfstudio.field_components.field_heads import (
-    DensityFieldHead,
-    FieldHead,
-    FieldHeadNames,
-    PredNormalsFieldHead,
-    RGBFieldHead,
-)
-from nerfstudio.field_components.mlp import MLP
-from nerfstudio.field_components.spatial_distortions import (
-    SceneContraction,
-    SpatialDistortion,
-)
+from nerfstudio.field_components.field_heads import FieldHeadNames, PredNormalsFieldHead
+from nerfstudio.field_components.spatial_distortions import SpatialDistortion
 from nerfstudio.fields.base_field import Field
 from nerfstudio.utils import math
 

nerfstudio/generative/stable_diffusion.py

@@ -1,3 +1,17 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 # From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py
 
 import matplotlib.pyplot as plt
@@ -91,8 +105,7 @@ class StableDiffusion(nn.Module):
         return text_embeddings
 
     def sds_loss(self, text_embeddings, nerf_output, guidance_scale=100):
-
-        nerf_output_np = nerf_output.view(3, 64, 64).permute(1, 2, 0).detach().cpu().numpy()
+        nerf_output_np = nerf_output.squeeze(0).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")
@@ -201,7 +214,7 @@ class StableDiffusion(nn.Module):
             latents=latents,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
-        )  # [1, 4, 64, 64]
+        )  # [1, 4, resolution, resolution]
 
         diffused_img = self.latents_to_img(latents)
         diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy()

nerfstudio/model_components/scene_colliders.py

@@ -107,6 +107,7 @@ class AABBBoxCollider(SceneCollider):
         ray_bundle.fars = fars[..., None]
         return ray_bundle
 
+
 class SphereCollider(SceneCollider):
     """Module for colliding rays with the scene box to compute near and far values.
 
@@ -121,9 +122,7 @@ class SphereCollider(SceneCollider):
         self.radius = radius
         self.near_plane = near_plane
 
-    def _intersect_with_sphere(
-        self, rays_o: TensorType["num_rays":..., 3], rays_d: TensorType["num_rays":..., 3]
-    ):
+    def _intersect_with_sphere(self, rays_o: TensorType["num_rays":..., 3], rays_d: TensorType["num_rays":..., 3]):
         """Returns collection of valid rays within a specified near/far bounding box along with a mask
         specifying which rays are valid
 
@@ -132,12 +131,12 @@ class SphereCollider(SceneCollider):
             rays_d: (num_rays, 3) ray directions
         """
 
-        a = (rays_d * rays_d).sum(dim = -1, keepdim = True)
+        a = (rays_d * rays_d).sum(dim=-1, keepdim=True)
         self.center = self.center.to(rays_o)
         b = 2 * (rays_o - self.center) * rays_d
-        b = b.sum(dim = -1, keepdim = True)
+        b = b.sum(dim=-1, keepdim=True)
         c = (rays_o - self.center) * (rays_o - self.center)
-        c = c.sum(dim = -1, keepdim = True) - self.radius ** 2
+        c = c.sum(dim=-1, keepdim=True) - self.radius**2
 
         # clamp to near plane
         near_plane = self.near_plane if self.training else 0

nerfstudio/models/dreamfusion.py

@@ -43,7 +43,7 @@ from nerfstudio.model_components.losses import (
     orientation_loss,
     pred_normal_loss,
 )
-from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler
+from nerfstudio.model_components.ray_samplers import UniformSampler
 from nerfstudio.model_components.renderers import (
     AccumulationRenderer,
     DepthRenderer,
@@ -76,6 +76,14 @@ class DreamFusionModelConfig(ModelConfig):
     """Initial strength of center density"""
     sphere_collider: bool = True
     """Use spherical collider instead of box"""
+    target_transmittance_start: float = 0.8
+    """target transmittance for opacity penalty. This is the percent of the scene that is 
+    background when rendered at the start of training"""
+    target_transmittance_end: float = 0.4
+    """target transmittance for opacity penalty. This is the percent of the scene that is
+    background when rendered at the end of training"""
+    transmittance_end_schedule: int = 1000
+    """number of iterations to reach target_transmittance_end"""
 
 
 class DreamFusionModel(Model):
@@ -97,6 +105,7 @@ class DreamFusionModel(Model):
         self.train_normals = False
         self.train_shaded = False
         self.density_strength = 1.0
+        self.target_transmittance = config.target_transmittance_start
         super().__init__(config=config, **kwargs)
 
     def get_training_callbacks(
@@ -119,6 +128,14 @@ class DreamFusionModel(Model):
         ):
             self.train_shaded = True
 
+        def update_target_transmittance(
+            self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
+        ):
+            if step < self.config.transmittance_end_schedule:
+                self.target_transmittance -= (1 / self.config.transmittance_end_schedule) * (
+                    self.config.target_transmittance_start - self.config.target_transmittance_end
+                )
+
         callbacks = [
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
@@ -138,6 +155,12 @@ class DreamFusionModel(Model):
                 func=start_shaded_training,
                 args=[self, training_callback_attributes],
             ),
+            TrainingCallback(
+                where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
+                update_every_num_iters=1,
+                func=update_target_transmittance,
+                args=[self, training_callback_attributes],
+            ),
         ]
         return callbacks
 
@@ -259,6 +282,9 @@ class DreamFusionModel(Model):
             field_outputs[FieldHeadNames.PRED_NORMALS],
         )
 
+        assert weights.shape[-1] == 1
+        outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01)
+
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
@@ -276,6 +302,17 @@ class DreamFusionModel(Model):
         loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
             outputs["rendered_pred_normal_loss"]
         )
+
+        loss_dict["alphas_loss"] = outputs["alphas_loss"].mean()
+
+        assert (
+            self.target_transmittance >= self.config.target_transmittance_end
+            and self.target_transmittance <= self.config.target_transmittance_start
+        )
+        loss_dict["opacity_loss"] = -torch.minimum(
+            (1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance)
+        )
+
         return loss_dict
 
     def get_image_metrics_and_images(

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -1,18 +1,32 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dreamfusion Pipeline and trainer"""
+
 from dataclasses import dataclass, field
-from typing import Type
+from typing import Optional, Type
 
 import matplotlib.pyplot as plt
 import numpy as np
-from torch.cuda.amp.grad_scaler import GradScaler
+import torch
 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
@@ -34,17 +48,25 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     """appended to prompt for front view"""
     back_prompt: str = ", back view"
     """appended to prompt for back view"""
-
+    alphas_penalty: bool = True
+    """enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
+    alphas_scale: float = 1
+    """scale for alphas penalty"""
     opacity_penalty: bool = True
-    """enables penalty to encourage transparent scenes"""
-    target_transmittance: float = 0.5
-    """target transmittance for opacity penalty"""
-
+    """enables penalty to encourage transparent scenes, as in "dreamfields" paper"""
+    opacity_scale: float = 1
+    """scale for opacity penalty"""
     guidance_scale: float = 100
     """guidance scale for sds loss"""
 
+    stablediffusion_device: Optional[str] = None
+
 
 class DreamfusionPipeline(VanillaPipeline):
+    """Dreamfusion pipeline"""
+
+    config: DreamfusionPipelineConfig
+
     def __init__(
         self,
         config: DreamfusionPipelineConfig,
@@ -55,7 +77,12 @@ class DreamfusionPipeline(VanillaPipeline):
     ):
         super().__init__(config, device, test_mode, world_size, local_rank)
         self.generative = True
-        self.sd = StableDiffusion(device)
+        self.sd_device = (
+            torch.device(device)
+            if self.config.stablediffusion_device is None
+            else torch.device(self.config.stablediffusion_device)
+        )
+        self.sd = StableDiffusion(self.sd_device)
         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}", "")
@@ -63,32 +90,37 @@ class DreamfusionPipeline(VanillaPipeline):
         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):
-        """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.
+    def get_train_loss_dict(self, step: int):
 
-        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)
+        metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
+        loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
 
-        train_output = model_outputs["train_output"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
+        train_output = (
+            model_outputs["train_output"]
+            .view(1, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3)
+            .permute(0, 3, 1, 2)
+        )
 
-        accumulation = model_outputs["accumulation"].view(64, 64).detach().cpu().numpy()
+        accumulation = (
+            model_outputs["accumulation"]
+            .view(self.config.datamanager.train_resolution, self.config.datamanager.train_resolution)
+            .detach()
+            .cpu()
+            .numpy()
+        )
         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 = model_outputs["background"].view(res, res, 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 = model_outputs["shaded"].view(res, res, 3).detach().cpu().numpy()
         # shaded = np.clip(shaded, 0.0, 1.0)
         # plt.imsave("nerf_textureless.jpg", shaded)
 
@@ -99,33 +131,27 @@ class DreamfusionPipeline(VanillaPipeline):
                 text_embedding = self.side_text_embedding
             elif batch["central"] > 135 and batch["central"] <= 225:
                 text_embedding = self.back_text_embedding
-            elif batch["central"] > 225 and batch["central"] <= 315:
+            else:  # batch["central"] > 225 and batch["central"] <= 315:
                 text_embedding = self.side_text_embedding
         else:
             text_embedding = self.base_text_embedding
 
         sds_loss, latents, grad = self.sd.sds_loss(
-            text_embedding, train_output, guidance_scale=self.config.guidance_scale
+            text_embedding.to(self.sd_device),
+            train_output.to(self.sd_device),
+            guidance_scale=int(self.config.guidance_scale),
         )
+        loss_dict["sds_loss"] = sds_loss.to(self.device)
         # TODO: opacity penalty using transmittance, not accumultation
         if self.config.opacity_penalty:
             accum_mean = np.mean(1.0 - accumulation)
             sds_loss *= np.min((0.5, accum_mean))
 
-        grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True)
-        # optimizers.scheduler_step_all(step)
-        # optimizers.optimizer_scaler_step_all(grad_scaler)
-        # grad_scaler.update()
-
-        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
-
-        normals_loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"]
-        grad_scaler.scale(normals_loss).backward()  # type: ignore
-        optimizers.optimizer_scaler_step_all(grad_scaler)
-
-        grad_scaler.update()
-        optimizers.scheduler_step_all(step)
+        model_outputs["latents"] = latents
+        model_outputs["grad"] = grad
 
         return model_outputs, loss_dict, metrics_dict
+
+    def forward(self):
+        """Not implemented since we only want the parameter saving of the nn module, but not forward()"""
+        raise NotImplementedError

nerfstudio/utils/math.py

@@ -195,4 +195,5 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
 
 # From https://github.com/ashawkey/stable-dreamfusion
 def safe_normalize(x, eps=1e-20):
+    """Safe normalization w/o divide by zero issues"""
     return x / torch.sqrt(torch.clamp(torch.sum(x * x, -1, keepdim=True), min=eps))

tests/test_train.py

@@ -15,15 +15,7 @@ from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConf
 from nerfstudio.engine.trainer import TrainerConfig
 from scripts.train import train_loop
 
-BLACKLIST = [
-    "base",
-    "semantic-nerfw",
-    "instant-ngp",
-    "instant-ngp-bounded",
-    "nerfacto",
-    "phototourism",
-    "depth-nerfacto",
-]
+BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "instant-ngp-bounded", "nerfacto", "phototourism", "dreamfusion"]
 
 
 def set_reduced_config(config: TrainerConfig):