diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 14ed3c8a2..3e6f3f80d 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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 diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 5dee76ceb..4dd44ceec 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/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]) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 78bf96f43..b52c80ad4 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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.""" diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 288dbc6fc..1bc12bef1 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/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)