Merge branch 'generative' of https://github.com/plenoptix/nerfactory into generative

nerfstudio/configs/method_configs.py

@@ -43,6 +43,7 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import (
 )
 from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
+from nerfstudio.models.dreamfusion import DreamFusionModelConfig
 from nerfstudio.models.instant_ngp import InstantNGPModelConfig
 from nerfstudio.models.mipnerf import MipNerfModel
 from nerfstudio.models.nerfacto import NerfactoModelConfig
@@ -62,6 +63,7 @@ descriptions = {
     "tensorf": "tensorf",
     "dnerf": "Dynamic-NeRF model. (slow)",
     "phototourism": "Uses the Phototourism data.",
+    "dreamfusion": "NeRF Model used in implementation of stable dreamfusion.",
 }
 
 method_configs["nerfacto"] = Config(
@@ -249,6 +251,37 @@ method_configs["phototourism"] = Config(
     vis="viewer",
 )
 
+method_configs["dreamfusion"] = Config(
+    method_name="dreamfusion",
+    trainer=TrainerConfig(
+        steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
+    ),
+    pipeline=VanillaPipelineConfig(
+        datamanager=VanillaDataManagerConfig(
+            dataparser=NerfstudioDataParserConfig(),
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=DreamFusionModelConfig(),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
+
 AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[  # Don't show unparseable (fixed) arguments in helptext.
     tyro.conf.FlagConversionOff[
         tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions)

nerfstudio/fields/dreamfusion_field.py

@@ -70,9 +70,6 @@ class DreamFusionField(Field):
         geo_feat_dim: output geo feat dimensions
         num_layers_color: number of hidden layers for color network
         hidden_dim_color: dimension of hidden layers for color network
-        appearance_embedding_dim: dimension of appearance embedding
-        transient_embedding_dim: dimension of transient embedding
-        use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference
         spatial_distortion: spatial distortion to apply to the scene
     """
 
@@ -85,7 +82,6 @@ class DreamFusionField(Field):
         num_layers_color: int = 3,
         hidden_dim_color: int = 64,
         hidden_dim_transient: int = 64,
-        use_pred_normals: bool = False,
         spatial_distortion: Optional[SpatialDistortion] = None,
     ) -> None:
         super().__init__()
@@ -94,7 +90,6 @@ class DreamFusionField(Field):
         self.geo_feat_dim = geo_feat_dim
 
         self.spatial_distortion = spatial_distortion
-        self.use_pred_normals = use_pred_normals
 
         num_levels = 16
         max_res = 1024
@@ -103,14 +98,6 @@ class DreamFusionField(Field):
         features_per_level = 2
         growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
-        self.direction_encoding = tcnn.Encoding(
-            n_input_dims=3,
-            encoding_config={
-                "otype": "SphericalHarmonics",
-                "degree": 4,
-            },
-        )
-
         self.position_encoding = tcnn.Encoding(
             n_input_dims=3,
             encoding_config={"otype": "Frequency", "n_frequencies": 2},
@@ -137,19 +124,18 @@ class DreamFusionField(Field):
         )
 
         # predicted normals
-        if self.use_pred_normals:
-            self.mlp_pred_normals = tcnn.Network(
-                n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,
-                n_output_dims=hidden_dim_transient,
-                network_config={
-                    "otype": "FullyFusedMLP",
-                    "activation": "ReLU",
-                    "output_activation": "None",
-                    "n_neurons": 64,
-                    "n_hidden_layers": 2,
-                },
-            )
-            self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
+        self.mlp_pred_normals = tcnn.Network(
+            n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,
+            n_output_dims=hidden_dim_transient,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "None",
+                "n_neurons": 64,
+                "n_hidden_layers": 2,
+            },
+        )
+        self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
 
         self.mlp_head = tcnn.Network(
             n_input_dims=self.geo_feat_dim,
@@ -190,22 +176,19 @@ class DreamFusionField(Field):
         outputs = {}
 
         directions = get_normalized_directions(ray_samples.frustums.directions)
-        directions_flat = directions.view(-1, 3)
-        d = self.direction_encoding(directions_flat)
 
         outputs_shape = ray_samples.frustums.directions.shape[:-1]
 
         # predicted normals
-        if self.use_pred_normals:
-            positions = ray_samples.frustums.get_positions()
+        positions = ray_samples.frustums.get_positions()
 
-            positions_flat = self.position_encoding(positions.view(-1, 3))
-            pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
+        positions_flat = self.position_encoding(positions.view(-1, 3))
+        pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
 
-            x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)
-            outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)
+        x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)
+        outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)
 
-        h = [density_embedding.view(-1, self.geo_feat_dim)]
+        h = density_embedding.view(-1, self.geo_feat_dim)
 
         rgb = self.mlp_head(h).view(*outputs_shape, -1).to(directions)
         outputs.update({FieldHeadNames.RGB: rgb})

nerfstudio/generative/stable_diffusion.py

@@ -0,0 +1,196 @@
+
+# From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py
+
+import torch
+from diffusers import AutoencoderKL, PNDMScheduler, UNet2DConditionModel
+from rich.console import Console
+from torch import nn
+from transformers import CLIPTextModel, CLIPTokenizer, logging
+
+CONSOLE = Console(width=120)
+logging.set_verbosity_error()
+IMG_DIM = 512
+
+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.')
+        except FileNotFoundError as e:
+            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"
+
+        self.num_train_timesteps = 1000
+        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}.')
+
+        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}.')
+
+        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.alphas = self.scheduler.alphas_cumprod.to(self.device)
+
+        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')
+
+        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')
+
+        with torch.no_grad():
+            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
+
+        # Cat for final embeddings
+        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+        return text_embeddings
+
+
+    def train_step(self, text_embeddings, pred_nerf, guidance_scale=100):
+        
+        # Upscale NeRF render for Stable Diffusion
+        pred_nerf = F.interpolate(pred_nerf, (IMG_DIM, IMG_DIM), mode='bilinear')
+
+        t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
+        latents = self.imgs_to_latent(pred_nerf)
+
+        # 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)
+            # 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 (high scale from paper!)
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        # w(t), sigma_t^2
+        w = (1 - self.alphas[t])
+        # w = self.alphas[t] ** 0.5 * (1 - self.alphas[t])
+        grad = w * (noise_pred - noise)
+
+        # clip grad for stable training?
+        # grad = grad.clamp(-10, 10)
+        grad = torch.nan_to_num(grad)
+
+        # manually backward, since we omitted an item in grad and cannot simply autodiff.
+        # _t = time.time()
+        latents.backward(gradient=grad, retain_graph=True)
+        # torch.cuda.synchronize(); print(f'[TIME] guiding: backward {time.time() - _t:.4f}s')
+
+        return 0 # dummy loss value
+
+    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)
+
+        self.scheduler.set_timesteps(num_inference_steps)
+
+        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']
+
+                # 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']
+        
+        return latents
+
+    def latents_to_img(self, latents):
+
+        latents = 1 / 0.18215 * latents
+
+        with torch.no_grad():
+            imgs = self.auto_encoder.decode(latents).sample
+
+        imgs = (imgs / 2 + 0.5).clamp(0, 1)
+        
+        return imgs
+
+    def imgs_to_latent(self, imgs):
+        # imgs: [B, 3, H, W]
+
+        imgs = 2 * imgs - 1
+
+        posterior = self.auto_encoder.encode(imgs).latent_dist
+        latents = posterior.sample() * 0.18215
+
+        return latents
+
+    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]
+        
+        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')
+
+        return diffused_img
+
+def seed_everything(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+if __name__ == '__main__':
+
+    import argparse
+
+    import matplotlib.pyplot as plt
+
+    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)
+    opt = parser.parse_args()
+
+    seed_everything(opt.seed)
+
+    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])

nerfstudio/model_components/losses.py

@@ -182,7 +182,7 @@ def orientation_loss(
     """
     w = weights
     n = normals
-    v = viewdirs
+    v = viewdirs * -1
     n_dot_v = (n * v[..., None, :]).sum(axis=-1)
     return (w[..., 0] * torch.fmin(torch.zeros_like(n_dot_v), n_dot_v) ** 2).sum(dim=-1)
 

nerfstudio/models/dreamfusion.py

@@ -30,19 +30,24 @@ from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 
 from nerfstudio.cameras.rays import RayBundle
 
-from nerfstudio.field_components.encodings import NeRFEncoding, 
+# from nerfstudio.field_components.encodings import NeRFEncoding,
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.fields.dreamfusion_field import DreamFusionField
-from nerfstudio.model_components.losses import MSELoss
+from nerfstudio.model_components.losses import (
+    MSELoss,
+    orientation_loss,
+    pred_normal_loss,
+)
 from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler
 from nerfstudio.model_components.renderers import (
     AccumulationRenderer,
     DepthRenderer,
+    NormalsRenderer,
     RGBRenderer,
 )
 from nerfstudio.model_components.scene_colliders import AABBBoxCollider
 from nerfstudio.models.base_model import Model, ModelConfig
-from nerfstudio.utils import colormaps, colors, misc
+from nerfstudio.utils import colormaps, colors, math, misc
 
 
 @dataclass
@@ -54,6 +59,10 @@ class DreamFusionModelConfig(ModelConfig):
 
     num_samples: int = 256
     """Number of samples in field evaluation"""
+    orientation_loss_mult: float = 0.0001
+    """Orientation loss multipier on computed normals."""
+    pred_normal_loss_mult: float = 0.001
+    """Predicted normal loss multiplier."""
 
 
 class DreamFusionModel(Model):
@@ -63,6 +72,8 @@ class DreamFusionModel(Model):
         config: DreamFusion configuration to instantiate model
     """
 
+    config: DreamFusionModelConfig
+
     def __init__(
         self,
         config: DreamFusionModelConfig,
@@ -87,6 +98,7 @@ class DreamFusionModel(Model):
         self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer()
+        self.renderer_normals = NormalsRenderer()
 
         # losses
         self.rgb_loss = MSELoss()
@@ -108,14 +120,46 @@ class DreamFusionModel(Model):
     def get_outputs(self, ray_bundle: RayBundle):
         # uniform sampling
         ray_samples_uniform = self.sampler_uniform(ray_bundle)
-        field_outputs = self.field(ray_samples_uniform)
+        field_outputs = self.field(ray_samples_uniform, compute_normals=True)
         weights = ray_samples_uniform.get_weights(field_outputs[FieldHeadNames.DENSITY])
 
         accumulation = self.renderer_accumulation(weights)
         depth = self.renderer_depth(weights, ray_samples_uniform)
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+        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
+        light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals)
+        light_d = math.safe_normalize(light_d)
+
+        ratio = 0.1
+        lambertian = ratio + (1 - ratio) * (pred_normals @ light_d).clamp(min=0)  # [N,]
+
+        shaded = lambertian.unsqueeze(-1).repeat(1, 3)
+        shaded_albedo = rgb * lambertian.unsqueeze(-1)
+
+        outputs = {
+            "rgb": rgb,
+            "accumulation": accumulation,
+            "depth": depth,
+            "normals": normals,
+            "pred_normals": pred_normals,
+            "shaded": shaded,
+            "shaded_albedo": shaded_albedo,
+        }
+
+        if self.training:
+            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],
+            )
 
-        outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth}
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
@@ -128,6 +172,16 @@ class DreamFusionModel(Model):
 
         loss_dict = {"rgb_loss": rgb_loss}
         loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients)
+        if self.training:
+            # orientation loss for computed normals
+            loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean(
+                outputs["rendered_orientation_loss"]
+            )
+
+            # ground truth supervision for normals
+            loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
+                outputs["rendered_pred_normal_loss"]
+            )
         return loss_dict
 
     def get_image_metrics_and_images(

nerfstudio/models/nerfacto.py

@@ -99,7 +99,7 @@ class NerfactoModelConfig(ModelConfig):
     distortion_loss_mult: float = 0.002
     """Distortion loss multiplier."""
     orientation_loss_mult: float = 0.0001
-    """Orientation loss multipier on computed noramls."""
+    """Orientation loss multipier on computed normals."""
     pred_normal_loss_mult: float = 0.001
     """Predicted normal loss multiplier."""
     use_proposal_weight_anneal: bool = True

nerfstudio/utils/math.py

@@ -191,3 +191,8 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
     """
 
     return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
+
+
+# From https://github.com/ashawkey/stable-dreamfusion
+def safe_normalize(x, eps=1e-20):
+    return x / torch.sqrt(torch.clamp(torch.sum(x * x, -1, keepdim=True), min=eps))

pyproject.toml

@@ -41,6 +41,8 @@ dependencies = [
     "rich==12.5.1",
     "tensorboard==2.9.0",
     "torch==1.12.1",
+    "diffusers==0.9.0",
+    "transformers==4.24.0",
     "torchmetrics[image]>=0.9.3",
     "torchtyping>=0.1.4",
     "torchvision>=0.13.0",