diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index a58e30624..68e337071 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -54,6 +54,7 @@ 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 method_configs: Dict[str, Config] = {} @@ -261,7 +262,7 @@ method_configs["dreamfusion"] = Config( trainer=TrainerConfig( steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True ), - pipeline=VanillaPipelineConfig( + pipeline=DreamfusionPipelineConfig( datamanager=DreamFusionDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index c05756de0..5f78951d5 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -237,13 +237,12 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me self.eval_cameras, _, _ = random_train_pose( self.config.num_eval_angles, self.config.eval_resolution, - device=device, - # radius_mean=1, - # radius_std=0.1, - # central_rotation_range=[180, 180], - # vertical_rotation_range=[0, 0], - # focal_range=[1, 1], - # jitter_std=0.1, + device=self.device, + radius_mean=1.5, + radius_std=0, + focal_range=[1, 1], + vertical_rotation_range=[0, 20], + # jitter_std=0, ) self.train_dataset = TrivialDataset(self.eval_cameras) @@ -256,14 +255,16 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me self.train_count += 1 - if step > 2000: - cameras, _, _ = random_train_pose( - self.config.train_images_per_batch, self.config.train_resolution, device=self.device - ) + # TODO Reimplement when cameras are fully working + # if step > 2000: + # cameras, _, _ = random_train_pose( + # self.config.train_images_per_batch, self.config.train_resolution, device=self.device + # ) - ray_bundle = cameras.generate_rays(torch.tensor([[i] for i in range(self.config.train_images_per_batch)])) - return ray_bundle, {"initialization": False} + # ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() + # return ray_bundle, {"initialization": False} + # TODO below cameras, vertical_rotation, central_rotation = random_train_pose( self.config.train_images_per_batch, self.config.train_resolution, @@ -271,10 +272,13 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me radius_mean=1.5, radius_std=0, focal_range=[1, 1], - vertical_rotation_range=[-180, 180], - jitter_std=0, + # vertical_rotation_range=[-180, 180], + # jitter_std=0, ) - ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() + # ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() + + camera_idx = torch.randint(0, self.eval_cameras.shape[0], [1], dtype=torch.long, device=self.device) + ray_bundle = self.eval_cameras.generate_rays(camera_idx).flatten() return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation, "initialization": True} diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index ec5a574b7..75db23fc1 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -311,13 +311,18 @@ class Trainer: """ self.optimizers.zero_grad_all() cpu_or_cuda_str = self.device.split(":")[0] - with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): - _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) + + if hasattr(self.pipeline, 'custom_step'): + _, loss_dict, metrics_dict = self.pipeline.custom_step(step=step, grad_scaler=self.grad_scaler, optimizers=self.optimizers) loss = functools.reduce(torch.add, loss_dict.values()) - # self.grad_scaler.scale(loss).backward() # type: ignore - # self.optimizers.optimizer_scaler_step_all(self.grad_scaler) - # self.grad_scaler.update() - self.optimizers.scheduler_step_all(step) + else: + with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): + _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) + loss = functools.reduce(torch.add, loss_dict.values()) + self.grad_scaler.scale(loss).backward() # type: ignore + self.optimizers.optimizer_scaler_step_all(self.grad_scaler) + self.grad_scaler.update() + self.optimizers.scheduler_step_all(step) # Merging loss and metrics dict into a single output. return loss, loss_dict, metrics_dict diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 6392a271e..5c50b4e46 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -2,6 +2,7 @@ # From https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py import matplotlib.pyplot as plt +import numpy as np import torch import torch.nn.functional as F from diffusers import AutoencoderKL, PNDMScheduler, UNet2DConditionModel @@ -75,7 +76,9 @@ class StableDiffusion(nn.Module): def sds_loss(self, text_embeddings, nerf_output, guidance_scale=100): - plt.imsave('nerf_output.png', nerf_output.view(3, 64, 64).permute(1, 2, 0).detach().cpu().numpy()) + nerf_output_np = nerf_output.view(3, 64, 64).permute(1, 2, 0).detach().cpu().numpy() + 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') t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) @@ -102,13 +105,17 @@ class StableDiffusion(nn.Module): # grad = grad.clamp(-10, 10) grad = torch.nan_to_num(grad) + # noise_difference = self.latents_to_img(grad) + # plt.imsave('noise_difference.png', noise_difference.view(3, 512, 512).permute(1, 2, 0).detach().cpu().numpy()) + # manually backward, since we omitted an item in grad and cannot simply autodiff. # _t = time.time() - latents.backward(gradient=grad, retain_graph=True) - rgb_loss = torch.square(noise_pred - noise) + noise_loss = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise))) - return rgb_loss + # print('SDS LOSS:', noise_loss) + + return noise_loss, latents, grad def produce_latents(self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None): diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 0ca922e37..6fe801200 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -38,7 +38,6 @@ from nerfstudio.engine.callbacks import ( # from nerfstudio.field_components.encodings import NeRFEncoding, from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.dreamfusion_field import DreamFusionField -from nerfstudio.generative.stable_diffusion import StableDiffusion from nerfstudio.model_components.losses import ( MSELoss, orientation_loss, @@ -73,8 +72,6 @@ class DreamFusionModelConfig(ModelConfig): """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" - prompt: str = "A high quality photo of a pineaple." - """Text prompt to optimize scene for.""" class DreamFusionModel(Model): @@ -95,8 +92,6 @@ class DreamFusionModel(Model): self.initialize_density = config.initialize_density self.train_normals = False super().__init__(config=config, **kwargs) - self.sd = StableDiffusion(kwargs['device']) - self.text_embedding = self.sd.get_text_embeds(config.prompt, "") def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes @@ -172,7 +167,7 @@ class DreamFusionModel(Model): density = field_outputs[FieldHeadNames.DENSITY] if self.initialize_density: pos = ray_samples_uniform.frustums.get_positions() - density_blob = 10 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] + density_blob = 50 * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] density += density_blob weights = ray_samples_uniform.get_weights(density) @@ -224,16 +219,9 @@ class DreamFusionModel(Model): return outputs def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: - # TODO: modify to not take in data images, need to get diffusion loss here # Scaling metrics by coefficients to create the losses. - device = outputs["rgb"].device - # change to key 'shaded_albedo' - albedo_output = outputs["rgb"].view(1, 64, 64, 3).permute(0, 3, 1, 2) - - rgb_loss = self.sd.sds_loss(self.text_embedding, albedo_output) - - loss_dict = {"rgb_loss": rgb_loss} + loss_dict = {} loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) if self.train_normals: # orientation loss for computed normals diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py new file mode 100644 index 000000000..f0c8cbd57 --- /dev/null +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass, field +from typing import Type + +from torch.cuda.amp.grad_scaler import GradScaler +from typing_extensions import Literal + +from nerfstudio.data.datamanagers.dreamfusion_datamanager import ( + DreamFusionDataManagerConfig, +) +from nerfstudio.engine.optimizers import Optimizers +from nerfstudio.generative.stable_diffusion import StableDiffusion +from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig +from nerfstudio.utils import profiler + + +@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""" + +class DreamfusionPipeline(VanillaPipeline): + + def __init__( + self, + config: DreamfusionPipelineConfig, + device: str, + test_mode: Literal["test", "val", "inference"] = "val", + world_size: int = 1, + local_rank: int = 0, + ): + super().__init__( + config, + device, + test_mode, + world_size, + local_rank + ) + self.sd = StableDiffusion(device) + self.text_embedding = self.sd.get_text_embeds("A high quality photo of a pineapple.", "") + + @profiler.time_function + def custom_step(self, step: int, grad_scaler: GradScaler, optimizers: Optimizers): + """This function gets your training loss dict. This will be responsible for + getting the next batch of data from the DataManager and interfacing with the + Model class, feeding the data to the model's forward function. + + Args: + step: current iteration step to update sampler if using DDP (distributed) + """ + 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) + + # Just uses albedo for now + albedo_output = model_outputs["rgb"].view(1, 64, 64, 3).permute(0, 3, 1, 2) + + sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output) + + grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True) + optimizers.optimizer_scaler_step_all(grad_scaler) + grad_scaler.update() + optimizers.scheduler_step_all(step) + + metrics_dict = self.model.get_metrics_dict(model_outputs, batch) + loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) + loss_dict["sds_loss"] = sds_loss + + return model_outputs, loss_dict, metrics_dict \ No newline at end of file