12 GB Dreamfusion!

nerfstudio/configs/method_configs.py

@@ -26,13 +26,10 @@ from nerfacc import ContractionType
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import ViewerConfig
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig
-<<<<<<< HEAD
+from nerfstudio.data.datamanagers.depth_datamanager import DepthDataManagerConfig
 from nerfstudio.data.datamanagers.dreamfusion_datamanager import (
     DreamFusionDataManagerConfig,
 )
-=======
-from nerfstudio.data.datamanagers.depth_datamanager import DepthDataManagerConfig
->>>>>>> 34b453a288df0ad4d5ce429d464b65199bcfea2b
 from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig
 from nerfstudio.data.datamanagers.variable_res_datamanager import (
     VariableResDataManagerConfig,
@@ -52,11 +49,8 @@ from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConf
 from nerfstudio.engine.schedulers import SchedulerConfig
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
-<<<<<<< HEAD
-from nerfstudio.models.dreamfusion import DreamFusionModelConfig
-=======
 from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
->>>>>>> 34b453a288df0ad4d5ce429d464b65199bcfea2b
+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

nerfstudio/generative/stable_diffusion.py

@@ -18,10 +18,10 @@ import matplotlib.pyplot as plt
 import numpy as np
 import torch
 import torch.nn.functional as F
-from diffusers import AutoencoderKL, PNDMScheduler, UNet2DConditionModel
+from diffusers import PNDMScheduler, StableDiffusionPipeline
 from rich.console import Console
 from torch import nn
-from transformers import CLIPTextModel, CLIPTokenizer, logging
+from transformers import logging
 
 CONSOLE = Console(width=120)
 logging.set_verbosity_error()
@@ -46,27 +46,13 @@ class StableDiffusion(nn.Module):
                 f"Cannot read hugging face token, make sure to save your Hugging Face token at `./HF_TOKEN`"
             )
 
-        SD_SOURCE = "runwayml/stable-diffusion-v1-5"
+        SD_SOURCE = "stabilityai/stable-diffusion-2-1"
         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,
@@ -75,6 +61,25 @@ class StableDiffusion(nn.Module):
         )
         self.alphas = self.scheduler.alphas_cumprod.to(self.device)
 
+        pipe = StableDiffusionPipeline.from_pretrained(SD_SOURCE, torch_dtype=torch.float16)
+        pipe = pipe.to(self.device)
+
+        # MEMORY IMPROVEMENTS
+        pipe.enable_attention_slicing()
+
+        # Needs xformers package (installation is from source https://github.com/facebookresearch/xformers)
+        # pipe.enable_xformers_memory_efficient_attention()
+
+        # More memory savings for 1/3rd performance
+        # pipe.enable_sequential_cpu_offload()
+
+        self.tokenizer = pipe.tokenizer
+        self.unet = pipe.unet
+        self.text_encoder = pipe.text_encoder
+        self.auto_encoder = pipe.vae
+
+        self.unet.to(memory_format=torch.channels_last)
+
         CONSOLE.print(f"Stable Diffusion loaded!")
 
     def get_text_embeds(self, prompt, negative_prompt):
@@ -109,9 +114,6 @@ class StableDiffusion(nn.Module):
         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")
-
-        # print(torch.max(nerf_output), torch.min(nerf_output))
-
         t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
         latents = self.imgs_to_latent(nerf_output)
         # predict the noise residual with unet, NO grad!
@@ -243,8 +245,10 @@ if __name__ == "__main__":
 
     device = torch.device("cuda")
 
-    sd = StableDiffusion(device)
+    with torch.no_grad():
 
-    imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.steps)
+        sd = StableDiffusion(device)
 
-    plt.imsave("test_sd.png", imgs[0])
+        imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.steps)
+
+        plt.imsave("test_sd.png", imgs[0])

nerfstudio/models/dreamfusion.py

@@ -62,7 +62,7 @@ class DreamFusionModelConfig(ModelConfig):
     _target: Type = field(default_factory=lambda: DreamFusionModel)
     """target class to instantiate"""
 
-    num_samples: int = 256
+    num_samples: int = 128
     """Number of samples in field evaluation"""
     orientation_loss_mult: float = 0.01
     """Orientation loss multipier on computed normals."""

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -95,6 +95,7 @@ class DreamfusionPipeline(VanillaPipeline):
         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)
@@ -136,12 +137,13 @@ class DreamfusionPipeline(VanillaPipeline):
         else:
             text_embedding = self.base_text_embedding
 
-        sds_loss, latents, grad = self.sd.sds_loss(
-            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)
+        with torch.autocast(device_type='cuda', dtype=torch.float16):
+            sds_loss, latents, grad = self.sd.sds_loss(
+                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)