Merge branch 'generative' of https://github.com/nerfstudio-project/nerfstudio into generative

nerfstudio/configs/method_configs.py

@@ -60,7 +60,6 @@ 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
 from nerfstudio.plugins.registry import discover_methods
 
@@ -344,7 +343,8 @@ method_configs["dreamfusion"] = TrainerConfig(
     steps_per_save=200,
     max_num_iterations=30000,
     mixed_precision=True,
-    pipeline=DreamfusionPipelineConfig(
+    pipeline=VanillaPipelineConfig(
+        generative=True,
         datamanager=DreamFusionDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,
@@ -352,22 +352,22 @@ method_configs["dreamfusion"] = TrainerConfig(
         ),
         model=DreamFusionModelConfig(
             eval_num_rays_per_chunk=1 << 15,
-            distortion_loss_mult=10.0,
+            distortion_loss_mult=0.02,
+            interlevel_loss_mult=1.0,
             orientation_loss_mult=0.1,
             max_res=256,
             sphere_collider=True,
             initialize_density=False,
             random_background=True,
-            interlevel_loss_mult=1.0,
-            proposal_warmup=500,
-            proposal_update_every=5,
+            proposal_warmup=1000,
+            proposal_update_every=2,
             proposal_weights_anneal_max_num_iters=100,
             start_lambertian_training=1000,
             start_normals_training=500,
             opacity_loss_mult=0.001,
+            prompting_type="location_based",
+            guidance_scale=100,
         ),
-        interpolated_prompting=False,
-        guidance_scale=100,
     ),
     optimizers={
         "proposal_networks": {

nerfstudio/generative/stable_diffusion_utils.py

@@ -0,0 +1,61 @@
+# 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.
+
+"""
+Stable diffusion utils
+"""
+
+def get_text_embedding(batch, prompting_type, text_embeddings, sd_device):
+    """"""
+
+    base = text_embeddings["base_text_embedding"]
+    front = text_embeddings["front_text_embedding"]
+    side = text_embeddings["side_text_embedding"]
+    top = text_embeddings["top_text_embedding"]
+    back = text_embeddings["back_text_embedding"]
+
+    assert prompting_type in ("location_based", "interpolated", "base")
+
+    if prompting_type == "location_based":
+        if batch["vertical"] < 40:
+            text_embedding = top
+        elif batch["central"] > 315 or batch["central"] <= 45:
+            text_embedding = front
+        elif batch["central"] > 45 and batch["central"] <= 135:
+            text_embedding = side
+        elif batch["central"] > 135 and batch["central"] <= 225:
+            text_embedding = back
+        else:  # batch["central"] > 225 and batch["central"] <= 315:
+            text_embedding = side
+
+    elif prompting_type == "interpolated":
+        horiz = batch["central"].to(sd_device)
+        vert = max(batch["vertical"].to(sd_device), 0)
+
+        if batch["central"] > 0 or batch["central"] <= 90:
+            text_embedding = (horiz) * side + (90 - horiz) * front
+        elif batch["central"] > 90 and batch["central"] <= 180:
+            text_embedding = (horiz - 90) * back + (180 - horiz) * side
+        elif batch["central"] > 180 and batch["central"] <= 270:
+            text_embedding = (horiz - 180) * side + (270 - horiz) * back
+        else:  # batch["central"] > 270 and batch["central"] <= 360:
+            text_embedding = (horiz - 270) * front + (360 - horiz) * side
+            
+        text_embedding = text_embedding / 90.0
+        text_embedding = (vert * text_embedding + (90 - vert) * top) / 90.0
+
+    elif prompting_type == "base":
+        text_embedding = base
+
+    return text_embedding

nerfstudio/model_components/ray_samplers.py

@@ -510,7 +510,7 @@ class VolumetricSampler(Sampler):
 
 class ProposalNetworkSampler(Sampler):
     """Sampler that uses a proposal network to generate samples.
-    
+
     Args:
         num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step.
         num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model.
@@ -518,6 +518,7 @@ class ProposalNetworkSampler(Sampler):
         single_jitter: Use a same random jitter for all samples along a ray.
         update_sched: A function that takes the iteration number of steps between updates.
         initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set.
+        pdf_histogram_padding: Padding to add to the histogram of the pdf.
     """
 
     def __init__(
@@ -528,6 +529,7 @@ class ProposalNetworkSampler(Sampler):
         single_jitter: bool = False,
         update_sched: Callable = lambda x: 1,
         initial_sampler: Optional[Sampler] = None,
+        pdf_histogram_padding: float = 0.01,
     ) -> None:
         super().__init__()
         self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
@@ -542,7 +544,9 @@ class ProposalNetworkSampler(Sampler):
             self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
         else:
             self.initial_sampler = initial_sampler
-        self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)
+        self.pdf_sampler = PDFSampler(
+            include_original=False, single_jitter=single_jitter, histogram_padding=pdf_histogram_padding
+        )
 
         self._anneal = 1.0
         self._steps_since_update = 0

nerfstudio/models/dreamfusion.py

@@ -19,11 +19,12 @@ DreamFusion implementation.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Dict, List, Tuple, Type
+from typing import Dict, List, Tuple, Type, Optional
 
 import numpy as np
 import torch
 from torch.nn import Parameter
+from torch.cuda.amp.grad_scaler import GradScaler
 
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.engine.callbacks import (
@@ -31,6 +32,8 @@ from nerfstudio.engine.callbacks import (
     TrainingCallbackAttributes,
     TrainingCallbackLocation,
 )
+from nerfstudio.generative.stable_diffusion import StableDiffusion
+from nerfstudio.generative.stable_diffusion_utils import get_text_embedding
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.fields.density_fields import HashMLPDensityField
 from nerfstudio.fields.dreamfusion_field import DreamFusionField
@@ -128,6 +131,25 @@ class DreamFusionModelConfig(ModelConfig):
     max_res: int = 256
     """Maximum resolution of the density field."""
 
+    location_based_prompting: bool = True
+    """enables location based prompting"""
+    interpolated_prompting: bool = False
+    """enables interpolated location prompting"""
+    prompting_type: str = "location_based"
+    """choose between location_based, interpolated, base"""
+    top_prompt: str = ", overhead view"
+    """appended to prompt for overhead view"""
+    side_prompt: str = ", side view"
+    """appended to prompt for side view"""
+    front_prompt: str = ", front view"
+    """appended to prompt for front view"""
+    back_prompt: str = ", back view"
+    """appended to prompt for back view"""
+    guidance_scale: float = 100
+    """guidance scale for sds loss"""
+    stablediffusion_device: Optional[str] = None
+    """device for stable diffusion"""
+    sd_version: str = '1-5' 
 
 class DreamFusionModel(Model):
     """DreamFusionModel Model
@@ -144,18 +166,44 @@ class DreamFusionModel(Model):
         **kwargs,
     ) -> None:
         self.prompt = config.prompt
+        self.cur_prompt = config.prompt
+        self.sd_version = config.sd_version
         self.initialize_density = config.initialize_density
         self.train_normals = False
         self.train_shaded = False
         self.random_background = config.random_background
         self.density_strength = 1.0
         self.target_transmittance = config.target_transmittance_start
+        self.grad_scaler = kwargs['grad_scaler']
+
+        self.prompting_type = config.prompting_type
+        self.guidance_scale = config.guidance_scale
+        self.top_prompt = config.top_prompt
+        self.side_prompt = config.side_prompt
+        self.back_prompt = config.back_prompt 
+        self.front_prompt = config.front_prompt
+
+        self.sd_device = (
+            torch.device(kwargs['device'])
+            if config.stablediffusion_device is None
+            else torch.device(config.stablediffusion_device)
+        )
+
         super().__init__(config=config, **kwargs)
 
     def populate_modules(self):
         """Set the fields and modules"""
         super().populate_modules()
 
+        self.sd = StableDiffusion(self.sd_device, version=self.sd_version)
+        self.text_embeddings = {
+            "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""),
+            "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""),
+            "side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""),
+            "back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""),
+            "base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "")
+        }
+
         # setting up fields
         self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res)
 
@@ -186,6 +234,7 @@ class DreamFusionModel(Model):
             single_jitter=self.config.use_single_jitter,
             update_sched=update_schedule,
             initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter),
+            pdf_histogram_padding=0.001,
         )
 
         # renderers
@@ -350,7 +399,7 @@ class DreamFusionModel(Model):
             shading_weight = 0.0
 
         shaded, shaded_albedo = self.shader_lambertian(
-            rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True
+            rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False
         )
 
         outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
@@ -412,6 +461,33 @@ class DreamFusionModel(Model):
         if self.config.opacity_penalty:
             loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean()
 
+        if self.prompt != self.cur_prompt:
+            self.cur_prompt = self.prompt
+            self.text_embeddings = {
+                "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""),
+                "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""),
+                "side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""),
+                "back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""),
+                "base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "")
+            }
+        
+        text_embedding = get_text_embedding(batch, self.prompting_type, self.text_embeddings, self.sd_device)
+            
+        train_output = (
+            outputs["train_output"].view(1, int(outputs["train_output"].shape[0] ** 0.5), int(outputs["train_output"].shape[0] ** 0.5), 3)
+            .permute(0, 3, 1, 2)
+        )
+
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            sds_loss = self.sd.sds_loss(
+                text_embedding.to(self.sd_device),
+                train_output.to(self.sd_device),
+                guidance_scale=int(self.guidance_scale),
+                grad_scaler=self.grad_scaler,
+            )
+
+        loss_dict["sds_loss"] = sds_loss.to(self.device)
+
         if self.training:
             loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss(
                 outputs["weights_list"], outputs["ray_samples_list"]

nerfstudio/pipelines/base_pipeline.py

@@ -191,6 +191,8 @@ class VanillaPipelineConfig(cfg.InstantiateConfig):
     """specifies the datamanager config"""
     model: ModelConfig = ModelConfig()
     """specifies the model config"""
+    generative: bool = False
+    """specifies whether the pipeline is for generative models"""
 
 
 class VanillaPipeline(Pipeline):
@@ -225,7 +227,7 @@ class VanillaPipeline(Pipeline):
         self.datamanager: VanillaDataManager = config.datamanager.setup(
             device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank
         )
-        self.generative = False
+        self.generative = config.generative
         self.datamanager.to(device)
         # TODO(ethan): get rid of scene_bounds from the model
         assert self.datamanager.train_dataset is not None, "Missing input dataset"
@@ -235,6 +237,7 @@ class VanillaPipeline(Pipeline):
             num_train_data=len(self.datamanager.train_dataset),
             metadata=self.datamanager.train_dataset.metadata,
             device=device,
+            grad_scaler=grad_scaler
         )
         self.model.to(device)
 

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -1,179 +0,0 @@
-# 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 Optional, Type
-
-import torch
-from torch.cuda.amp.grad_scaler import GradScaler
-from typing_extensions import Literal
-
-from nerfstudio.data.datamanagers.dreamfusion_datamanager import (
-    DreamFusionDataManagerConfig,
-)
-from nerfstudio.generative.stable_diffusion import StableDiffusion
-from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig
-
-
-@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"""
-    # prompt: str = "A high-quality photo of a tree frog on a stump"
-    # """prompt for stable dreamfusion"""
-
-    location_based_prompting: bool = True
-    """enables location based prompting"""
-    interpolated_prompting: bool = False
-    """enables interpolated location prompting"""
-    top_prompt: str = ", overhead view"
-    """appended to prompt for overhead view"""
-    side_prompt: str = ", side view"
-    """appended to prompt for side view"""
-    front_prompt: str = ", front view"
-    """appended to prompt for front view"""
-    back_prompt: str = ", back view"
-    """appended to prompt for back view"""
-    guidance_scale: float = 100
-    """guidance scale for sds loss"""
-    sd_version: str = "1-5"
-    """Stable Diffusion version for dreamfusion"""
-
-    stablediffusion_device: Optional[str] = None
-
-
-class DreamfusionPipeline(VanillaPipeline):
-    """Dreamfusion pipeline"""
-
-    config: DreamfusionPipelineConfig
-
-    def __init__(
-        self,
-        config: DreamfusionPipelineConfig,
-        device: str,
-        test_mode: Literal["test", "val", "inference"] = "val",
-        world_size: int = 1,
-        local_rank: int = 0,
-        grad_scaler: Optional[GradScaler] = None,
-    ):
-        super().__init__(config, device, test_mode, world_size, local_rank, grad_scaler)
-        self.generative = True
-        self.cur_prompt = self.model.prompt
-        self.grad_scaler = grad_scaler
-        self.sd_device = (
-            torch.device(device)
-            if self.config.stablediffusion_device is None
-            else torch.device(self.config.stablediffusion_device)
-        )
-        self.location_based_prompting = config.location_based_prompting
-        self.top_prompt = config.top_prompt
-        self.side_prompt = config.side_prompt
-        self.back_prompt = config.back_prompt 
-        self.front_prompt = config.front_prompt
-        self.sd = StableDiffusion(self.sd_device, version=config.sd_version)
-        
-        if config.location_based_prompting:
-            self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "")
-            self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "")
-            self.side_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", "")
-            self.back_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", "")
-        else:
-            self.base_text_embedding = self.sd.get_text_embeds(self.cur_prompt, "")
-
-    def get_train_loss_dict(self, step: int):
-
-        if self.world_size > 1 and step:
-            assert self.datamanager.train_sampler is not None
-            self.datamanager.train_sampler.set_epoch(step)
-
-        if self.model.prompt != self.cur_prompt:
-            self.cur_prompt = self.model.prompt
-            if self.location_based_prompting:
-                self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "")
-                self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "")
-                self.side_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", "")
-                self.back_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", "")
-            else:
-                self.base_text_embedding = self.sd.get_text_embeds(self.cur_prompt, "")
-
-        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, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3)
-            .permute(0, 3, 1, 2)
-        )
-
-        if self.config.location_based_prompting:
-            if self.config.interpolated_prompting:
-                horiz = batch["central"].to(self.sd_device)
-                vert = max(batch["vertical"].to(self.sd_device), 0)
-
-                if batch["central"] > 0 or batch["central"] <= 90:
-                    text_embedding = (
-                        (horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding
-                    ) / 90.0
-                elif batch["central"] > 90 and batch["central"] <= 180:
-                    text_embedding = (
-                        (horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding
-                    ) / 90.0
-                elif batch["central"] > 180 and batch["central"] <= 270:
-                    text_embedding = (
-                        (horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding
-                    ) / 90.0
-                else:  # batch["central"] > 270 and batch["central"] <= 360:
-                    text_embedding = (
-                        (horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding
-                    ) / 90.0
-
-                text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0
-
-            else:
-                if batch["vertical"] < 40:
-                    text_embedding = self.top_text_embedding
-                elif batch["central"] > 315 or batch["central"] <= 45:
-                    text_embedding = self.front_text_embedding
-                elif batch["central"] > 45 and batch["central"] <= 135:
-                    text_embedding = self.side_text_embedding
-                elif batch["central"] > 135 and batch["central"] <= 225:
-                    text_embedding = self.back_text_embedding
-                else:  # batch["central"] > 225 and batch["central"] <= 315:
-                    text_embedding = self.side_text_embedding
-        else:
-            text_embedding = self.base_text_embedding
-
-        with torch.autocast(device_type="cuda", dtype=torch.float16):
-            sds_loss = self.sd.sds_loss(
-                text_embedding.to(self.sd_device),
-                train_output.to(self.sd_device),
-                guidance_scale=int(self.config.guidance_scale),
-                grad_scaler=self.grad_scaler,
-            )
-
-        loss_dict["sds_loss"] = sds_loss.to(self.device)
-
-        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