adding dream embedding

.vscode/settings.json

@@ -1,125 +0,0 @@
-{
-  "files.watcherExclude": {
-    "**/__pycache__/": true,
-    "**/.ipynb_checkpoints": true,
-    "**/*.ipynb": true,
-    "**/old*": true,
-    "**/runs*": true,
-    "**/temp*": true
-  },
-  "editor.defaultFormatter": "esbenp.prettier-vscode",
-  "[javascript]": {
-    "editor.defaultFormatter": "esbenp.prettier-vscode"
-  },
-  "[javascriptreact]": {
-    "editor.defaultFormatter": "esbenp.prettier-vscode"
-  },
-  "[typescript]": {
-    "editor.defaultFormatter": "esbenp.prettier-vscode"
-  },
-  "[typescriptreact]": {
-    "editor.defaultFormatter": "esbenp.prettier-vscode"
-  },
-  "prettier.configPath": ".prettierrc.js",
-  "typescript.suggestionActions.enabled": false,
-  "javascript.suggestionActions.enabled": false,
-  "[python]": {
-    "editor.defaultFormatter": "ms-python.python",
-    "editor.codeActionsOnSave": {
-      "source.organizeImports": true
-    }
-  },
-  "editor.formatOnSave": true,
-  "python.envFile": "${workspaceFolder}/.env",
-  "python.formatting.provider": "black",
-  "python.formatting.blackArgs": ["--line-length=120"],
-  "python.linting.pylintEnabled": true,
-  "python.linting.flake8Enabled": false,
-  "python.linting.enabled": true,
-  "python.sortImports.args": ["--src=${workspaceFolder}", "--profile", "black"],
-  "python.testing.unittestEnabled": false,
-  "python.testing.pytestEnabled": true,
-  "python.linting.ignorePatterns": [
-    "**/site-packages/**/*.py",
-    ".vscode/*.py",
-    "external/**/*.py"
-  ],
-  "esbonio.sphinx.confDir": "",
-  // eslint extension options
-  "javascript.validate.enable": false,
-  "eslint.enable": true,
-  "eslint.validate": [
-    "javascript",
-    "javascriptreact",
-    "typescript",
-    "typescriptreact"
-  ],
-  "C_Cpp.errorSquiggles": "Disabled",
-  "files.associations": {
-    "array": "cpp",
-    "bitset": "cpp",
-    "string_view": "cpp",
-    "initializer_list": "cpp",
-    "utility": "cpp",
-    "__hash_table": "cpp",
-    "__split_buffer": "cpp",
-    "deque": "cpp",
-    "iterator": "cpp",
-    "queue": "cpp",
-    "stack": "cpp",
-    "string": "cpp",
-    "unordered_map": "cpp",
-    "vector": "cpp",
-    "atomic": "cpp",
-    "*.tcc": "cpp",
-    "cctype": "cpp",
-    "chrono": "cpp",
-    "clocale": "cpp",
-    "cmath": "cpp",
-    "codecvt": "cpp",
-    "condition_variable": "cpp",
-    "cstdarg": "cpp",
-    "cstddef": "cpp",
-    "cstdint": "cpp",
-    "cstdio": "cpp",
-    "cstdlib": "cpp",
-    "cstring": "cpp",
-    "ctime": "cpp",
-    "cwchar": "cpp",
-    "cwctype": "cpp",
-    "exception": "cpp",
-    "algorithm": "cpp",
-    "filesystem": "cpp",
-    "functional": "cpp",
-    "memory": "cpp",
-    "memory_resource": "cpp",
-    "optional": "cpp",
-    "ratio": "cpp",
-    "system_error": "cpp",
-    "tuple": "cpp",
-    "type_traits": "cpp",
-    "fstream": "cpp",
-    "iomanip": "cpp",
-    "iosfwd": "cpp",
-    "iostream": "cpp",
-    "istream": "cpp",
-    "limits": "cpp",
-    "mutex": "cpp",
-    "new": "cpp",
-    "ostream": "cpp",
-    "sstream": "cpp",
-    "stdexcept": "cpp",
-    "streambuf": "cpp",
-    "thread": "cpp",
-    "typeinfo": "cpp",
-    "__nullptr": "cpp",
-    "__config": "cpp",
-    "__locale": "cpp",
-    "__bit_reference": "cpp",
-    "ios": "cpp",
-    "__atomic": "cpp",
-    "__node_handle": "cpp"
-  },
-  "python.analysis.typeCheckingMode": "basic",
-  "python.analysis.diagnosticMode": "workspace"
-}

nerfstudio/configs/method_configs.py

@@ -50,6 +50,7 @@ from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
 from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
+from nerfstudio.models.dreamembedding import DreamEmbeddingModelConfig
 from nerfstudio.models.dreamfusion import DreamFusionModelConfig
 from nerfstudio.models.instant_ngp import InstantNGPModelConfig
 from nerfstudio.models.mipnerf import MipNerfModel
@@ -428,6 +429,53 @@ method_configs["dreamfusion"] = TrainerConfig(
     vis="viewer",
 )
 
+method_configs["dreamembed"] = TrainerConfig(
+    method_name="dreamembed",
+    experiment_name="dream-embed",
+    steps_per_eval_batch=50,
+    steps_per_eval_image=50,
+    steps_per_save=200,
+    max_num_iterations=20000,
+    mixed_precision=True,
+    pipeline=VanillaPipelineConfig(
+        generative=True,
+        datamanager=DreamFusionDataManagerConfig(
+            horizontal_rotation_warmup=4000,
+        ),
+        model=DreamEmbeddingModelConfig(
+            eval_num_rays_per_chunk=1 << 15,
+            distortion_loss_mult=0.02,
+            interlevel_loss_mult=100.0,
+            orientation_loss_mult=0.1,
+            max_res=256,
+            sphere_collider=True,
+            initialize_density=False,
+            random_background=True,
+            proposal_warmup=500,
+            proposal_update_every=0,
+            proposal_weights_anneal_max_num_iters=100,
+            start_normals_training=1000,
+            opacity_loss_mult=0.001,
+            positional_prompting="discrete",
+            guidance_scale=100,
+        ),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(
+                warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear"
+            ),
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
 method_configs["nerfplayer-nerfacto"] = TrainerConfig(
     method_name="nerfplayer-nerfacto",
     pipeline=VanillaPipelineConfig(

nerfstudio/data/datamanagers/dreamfusion_datamanager.py

@@ -83,7 +83,7 @@ def random_train_pose(
         poses: [size, 4, 4]
     """
 
-    vertical_rotation_range = [vertical_rotation_range[0] + 90, vertical_rotation_range[1] + 90]
+    vertical_rotation_range = (vertical_rotation_range[0] + 90, vertical_rotation_range[1] + 90)
     # This is the uniform sample on the part of the sphere we care about where 0 = 0 degrees and 1 = 360 degrees
     sampled_uniform = (
         torch.rand(size) * (vertical_rotation_range[1] - vertical_rotation_range[0]) + vertical_rotation_range[0]
@@ -157,7 +157,7 @@ class DreamFusionDataManagerConfig(DataManagerConfig):
     """Number of images per batch for training"""
     eval_images_per_batch: int = 1
     """Number of images per batch for evaluation"""
-    radius_mean: float = 2.2
+    radius_mean: float = 3.0
     """Mean radius of camera orbit"""
     radius_std: float = 0.1
     """Std of radius of camera orbit"""
@@ -289,6 +289,23 @@ class DreamFusionDataManager(DataManager):  # pylint: disable=abstract-method
         ).flatten()
 
         return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation}
+    
+    def next_eval_image(self, step:int) -> Tuple[int, RayBundle, Dict]:
+        cameras, vertical_rotation, central_rotation = random_train_pose(
+            self.config.eval_images_per_batch,
+            self.config.eval_resolution,
+            device=self.device,
+            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,
+            center=self.config.center,
+        )
+        ray_bundle = cameras.generate_rays(
+            torch.tensor([[i] for i in range(self.config.train_images_per_batch)])
+        )
+        return 0, ray_bundle, {"vertical": vertical_rotation, "central": central_rotation}
 
     def get_train_rays_per_batch(self) -> int:
         return self.config.train_resolution**2

nerfstudio/generative/stable_diffusion.py

@@ -188,6 +188,42 @@ class StableDiffusion(nn.Module):
 
         return text_embeddings
 
+    def sds_loss_latent(
+        self,
+        text_embeddings: TensorType["N", "max_length", "embed_dim"],
+        latents: TensorType["BS", 3, "H", "W"],
+        guidance_scale: float = 100.0,
+        grad_scaler: Optional[GradScaler] = None,
+    ) -> torch.Tensor:
+        # Disable autocast when using diffusers pipeline.
+        with torch.autocast(device_type="cuda", enabled=False):
+            t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
+            latents = latents.half()
+            # predict the noise residual with unet, NO grad!
+            with torch.no_grad():
+                # add noise
+                noise = torch.randn_like(latents)
+                latents_noisy = self.scheduler.add_noise(latents, noise, t)  # type: ignore
+                # pred noise
+                latent_model_input = torch.cat([latents_noisy] * 2)
+                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_text + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+            # w(t), sigma_t^2
+            w = 1 - self.alphas[t]
+
+            grad = w * (noise_pred - noise)
+            grad = torch.nan_to_num(grad)
+
+            if grad_scaler is not None:
+                latents = grad_scaler.scale(latents)
+            loss = _SDSGradient.apply(latents, grad)
+
+        return loss
+
     def sds_loss(
         self,
         text_embeddings: TensorType["N", "max_length", "embed_dim"],
@@ -226,7 +262,7 @@ class StableDiffusion(nn.Module):
             # w(t), sigma_t^2
             w = 1 - self.alphas[t]
 
-            grad = w * (noise_pred - noise)
+            grad = w * (noise_pred - noise) * 0.01
             grad = torch.nan_to_num(grad)
 
             if grad_scaler is not None:

nerfstudio/models/dreamembedding.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Tuple, Type
 
+import mediapy
 import numpy as np
 import torch
 from torch.nn import Parameter
@@ -32,15 +33,15 @@ from nerfstudio.model_components.ray_samplers import (
     UniformSampler,
 )
 from nerfstudio.model_components.renderers import FeatureRenderer
-
 from nerfstudio.model_components.scene_colliders import AABBBoxCollider, SphereCollider
 from nerfstudio.model_components.shaders import LambertianShader, NormalsShader
-
 from nerfstudio.models.base_model import Model, ModelConfig
+from nerfstudio.models.dreamfusion import DreamFusionModel, DreamFusionModelConfig
 from nerfstudio.utils import colormaps, colors, math, misc
 
+
 @dataclass
-class DreamEmbeddingModelConfig(ModelConfig):
+class DreamEmbeddingModelConfig(DreamFusionModelConfig):
     """DreamFusion model config"""
 
     _target: Type = field(default_factory=lambda: DreamEmbeddingModel)
@@ -133,7 +134,7 @@ class DreamEmbeddingModelConfig(ModelConfig):
     sd_version: str = "1-5"
 
 
-class DreamEmbeddingModel(Model):
+class DreamEmbeddingModel(DreamFusionModel):
     """DreamEmbeddingModel Model
 
     Args:
@@ -157,19 +158,15 @@ class DreamEmbeddingModel(Model):
         # setting up fields
         self.field = DreamEmbeddingField(self.scene_box.aabb, max_res=self.config.max_res)
 
-        self.feature_renderer = FeatureRenderer()
-
+        self.renderer_feature = FeatureRenderer()
 
     def get_outputs(self, ray_bundle: RayBundle):  # pylint: disable=too-many-statements
         # uniform sampling
-        background_rgb = self.field.get_background_feature(ray_bundle)
+        background_feature = self.field.get_background_feature(ray_bundle)
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         field_outputs = self.field(ray_samples, compute_normals=True)
         density = field_outputs[FieldHeadNames.DENSITY]
 
-
-        # KEEP CHANIGNNG FROM HERE
-
         if self.initialize_density:
             pos = ray_samples.frustums.get_positions()
             density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None]
@@ -181,21 +178,32 @@ class DreamEmbeddingModel(Model):
 
         accumulation = self.renderer_accumulation(weights)
         depth = self.renderer_depth(weights, ray_samples)
-        rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+        features = self.renderer_feature(features=field_outputs[FieldHeadNames.FEATURE], weights=weights)
+        # rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
 
         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
 
-        background = accum_mask_inv * background_rgb
+        background = accum_mask_inv * background_feature
 
         outputs = {
-            "rgb_only": rgb,
-            "background_rgb": background_rgb,
+            "feature_only": features,
+            "background_feature": background_feature,
             "background": background,
             "accumulation": accum_mask,
             "depth": depth,
         }
 
+        latents = accum_mask * features + background
+        outputs["latents"] = latents
+
+        latents_input = latents.view(1, int(latents.shape[0] ** 0.5), int(latents.shape[0] ** 0.5), 4).half()
+        latents_input = latents_input.permute(0, 3, 1, 2)
+        rgb = self.sd.latents_to_img(latents_input).permute(0, 2, 3, 1)
+        rgb = rgb.view(512 * 512, 3)
+        rgb = rgb.to(torch.float32)
+        outputs["rgb"] = rgb
+
         # These use a lot of GPU memory, so we avoid storing them for eval.
         if self.training:
             outputs["weights_list"] = weights_list
@@ -205,55 +213,109 @@ class DreamEmbeddingModel(Model):
             outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
 
         normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
-        pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
-
-        # lambertian shading
-        if self.config.random_light_source:  # and self.training:
-            light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals)
-        else:
-            light_d = ray_bundle.origins[0]
-        light_d = math.safe_normalize(light_d)
-
-        if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training:
-            shading_weight = 0.9
-        else:
-            shading_weight = 0.0
-
-        shaded, shaded_albedo = self.shader_lambertian(
-            rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False
-        )
-
         outputs["normals"] = self.shader_normals(normals, weights=accum_mask)
-        outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask)
-        outputs["shaded"] = accum_mask * shaded
-        outputs["other_train_output"] = accum_mask * shaded_albedo + background
-        outputs["shaded_albedo"] = accum_mask * shaded_albedo
-        outputs["rgb"] = accum_mask * rgb + background
-
-        if shading_weight > 0:
-            samp = np.random.random_sample()
-            if samp > 0.5 and not self.training:
-                outputs["train_output"] = outputs["shaded"]
-            elif samp < 0.2 and self.random_background:
-                rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device)
-                outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv
-            else:
-                outputs["train_output"] = accum_mask * shaded_albedo + background
-        else:
-            outputs["train_output"] = outputs["rgb"]
 
         outputs["rendered_orientation_loss"] = orientation_loss(
             weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
         )
 
-        outputs["rendered_pred_normal_loss"] = pred_normal_loss(
-            weights.detach(),
-            field_outputs[FieldHeadNames.NORMALS].detach(),
-            field_outputs[FieldHeadNames.PRED_NORMALS],
-        )
-
         assert weights.shape[-1] == 1
         if self.config.opacity_penalty:
             outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult
 
-        return outputs
+        return outputs
+
+    def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
+        # Scaling metrics by coefficients to create the losses.
+
+        loss_dict = {}
+        loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)
+        if self.train_normals:
+            # orientation loss for computed normals
+            loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean(
+                outputs["rendered_orientation_loss"]
+            )
+        else:
+            loss_dict["orientation_loss"] = 0
+
+        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.update_prompt(
+                base_prompt=self.cur_prompt,
+                top_prompt=self.cur_prompt + self.top_prompt,
+                side_prompt=self.cur_prompt + self.side_prompt,
+                back_prompt=self.cur_prompt + self.back_prompt,
+                front_prompt=self.cur_prompt + self.front_prompt,
+            )
+
+        text_embedding = self.text_embeddings.get_text_embedding(
+            vertical_angle=batch["vertical"], horizontal_angle=batch["central"]
+        )
+
+        latents = (
+            outputs["latents"]
+            .view(1, int(outputs["latents"].shape[0] ** 0.5), int(outputs["latents"].shape[0] ** 0.5), 4)
+            .permute(0, 3, 1, 2)
+        )
+
+        # 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)
+        # )
+
+        # 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,
+        # )
+
+        sds_loss = self.sd.sds_loss_latent(
+            text_embedding.to(self.sd_device),
+            latents.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"]
+            )
+            loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
+                outputs["weights_list"], outputs["ray_samples_list"]
+            )
+        return loss_dict
+
+    def get_image_metrics_and_images(
+        self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
+        acc = colormaps.apply_colormap(outputs["accumulation"])
+        depth = colormaps.apply_depth_colormap(
+            outputs["depth"],
+            accumulation=outputs["accumulation"],
+        )
+        prop_depth_0 = colormaps.apply_depth_colormap(
+            outputs["prop_depth_0"],
+            accumulation=outputs["accumulation"],
+        )
+        prop_depth_1 = colormaps.apply_depth_colormap(
+            outputs["prop_depth_1"],
+            accumulation=outputs["accumulation"],
+        )
+
+        metrics_dict = {}
+        images_dict = {
+            "img": outputs["rgb"].reshape(512, 512, 3), # this is a bit weird
+            "accumulation": acc,
+            "depth": depth,
+            "prop_depth_0": prop_depth_0,
+            "prop_depth_1": prop_depth_1,
+            "normals": outputs["normals"],
+        }
+        return metrics_dict, images_dict

nerfstudio/utils/writer.py

@@ -283,7 +283,7 @@ class WandbWriter(Writer):
     """WandDB Writer Class"""
 
     def __init__(self, log_dir: Path):
-        wandb.init(project="nerfstudio-project", dir=str(log_dir), reinit=True)
+        wandb.init(project="dream-embed", dir=str(log_dir), reinit=True)
 
     def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None:
         image = torch.permute(image, (2, 0, 1))