From ca767bc9c7e3253aa323b20a0fe557cafc78893d Mon Sep 17 00:00:00 2001 From: terrance Date: Thu, 16 Feb 2023 23:03:43 +0000 Subject: [PATCH 1/4] remove dreamfusion pipeline file --- nerfstudio/configs/method_configs.py | 7 +- nerfstudio/models/dreamfusion.py | 76 ++++++++++++++- nerfstudio/pipelines/base_pipeline.py | 5 +- nerfstudio/pipelines/dreamfusion_pipeline.py | 98 ++++++++++---------- 4 files changed, 131 insertions(+), 55 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index d0c76262d..270c7598e 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -344,7 +344,8 @@ method_configs["dreamfusion"] = TrainerConfig( steps_per_save=200, max_num_iterations=30000, mixed_precision=True, - pipeline=DreamfusionPipelineConfig( + pipeline=VanillaPipelineConfig( + generative=True, datamanager=DreamFusionDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, @@ -365,9 +366,9 @@ method_configs["dreamfusion"] = TrainerConfig( start_lambertian_training=1000, start_normals_training=500, opacity_loss_mult=0.001, + prompting_type="location_based", + guidance_scale=100, ), - interpolated_prompting=False, - guidance_scale=100, ), optimizers={ "proposal_networks": { diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index c2a716946..c4e705664 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -19,11 +19,12 @@ DreamFusion implementation. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Tuple, Type +from typing import Dict, List, Tuple, Type, Optional import numpy as np import torch from torch.nn import Parameter +from torch.cuda.amp.grad_scaler import GradScaler from nerfstudio.cameras.rays import RayBundle from nerfstudio.engine.callbacks import ( @@ -31,6 +32,8 @@ from nerfstudio.engine.callbacks import ( TrainingCallbackAttributes, TrainingCallbackLocation, ) +from nerfstudio.generative.stable_diffusion import StableDiffusion +from nerfstudio.generative.stable_diffusion_utils import get_text_embedding from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.fields.dreamfusion_field import DreamFusionField @@ -128,6 +131,23 @@ class DreamFusionModelConfig(ModelConfig): max_res: int = 256 """Maximum resolution of the density field.""" + location_based_prompting: bool = True + """enables location based prompting""" + interpolated_prompting: bool = False + """enables interpolated location prompting""" + prompting_type: str = "location_based" + """choose between location_based, interpolated, base""" + top_prompt: str = ", overhead view" + """appended to prompt for overhead view""" + side_prompt: str = ", side view" + """appended to prompt for side view""" + front_prompt: str = ", front view" + """appended to prompt for front view""" + back_prompt: str = ", back view" + """appended to prompt for back view""" + guidance_scale: float = 100 + """guidance scale for sds loss""" + stablediffusion_device: Optional[str] = None class DreamFusionModel(Model): """DreamFusionModel Model @@ -144,18 +164,43 @@ class DreamFusionModel(Model): **kwargs, ) -> None: self.prompt = config.prompt + self.cur_prompt = config.prompt self.initialize_density = config.initialize_density self.train_normals = False self.train_shaded = False self.random_background = config.random_background self.density_strength = 1.0 self.target_transmittance = config.target_transmittance_start + self.grad_scaler = kwargs['grad_scaler'] + + self.prompting_type = config.prompting_type + self.guidance_scale = config.guidance_scale + self.top_prompt = config.top_prompt + self.side_prompt = config.side_prompt + self.back_prompt = config.back_prompt + self.front_prompt = config.front_prompt + + self.sd_device = ( + torch.device(kwargs['device']) + if config.stablediffusion_device is None + else torch.device(config.stablediffusion_device) + ) + super().__init__(config=config, **kwargs) def populate_modules(self): """Set the fields and modules""" super().populate_modules() + self.sd = StableDiffusion(self.sd_device) + self.text_embeddings = { + "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""), + "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""), + "side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""), + "back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""), + "base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "") + } + # setting up fields self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res) @@ -350,7 +395,7 @@ class DreamFusionModel(Model): shading_weight = 0.0 shaded, shaded_albedo = self.shader_lambertian( - rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True + rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False ) outputs["normals"] = self.shader_normals(normals, weights=accum_mask) @@ -412,6 +457,33 @@ class DreamFusionModel(Model): 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 = { + "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""), + "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""), + "side_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", ""), + "back_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", ""), + "base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, "") + } + + text_embedding = get_text_embedding(batch, self.prompting_type, self.text_embeddings, self.sd_device) + + 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) + ) + + with torch.autocast(device_type="cuda", dtype=torch.float16): + 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, + ) + + 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"] diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 3b7c45ecb..ac1b73dd6 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -191,6 +191,8 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """specifies the datamanager config""" model: ModelConfig = ModelConfig() """specifies the model config""" + generative: bool = False + """specifies whether the pipeline is for generative models""" class VanillaPipeline(Pipeline): @@ -225,7 +227,7 @@ class VanillaPipeline(Pipeline): self.datamanager: VanillaDataManager = config.datamanager.setup( device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank ) - self.generative = False + self.generative = config.generative self.datamanager.to(device) # TODO(ethan): get rid of scene_bounds from the model assert self.datamanager.train_dataset is not None, "Missing input dataset" @@ -235,6 +237,7 @@ class VanillaPipeline(Pipeline): num_train_data=len(self.datamanager.train_dataset), metadata=self.datamanager.train_dataset.metadata, device=device, + grad_scaler=grad_scaler ) self.model.to(device) diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 692a92a71..88ec2feb9 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -86,7 +86,7 @@ class DreamfusionPipeline(VanillaPipeline): self.back_prompt = config.back_prompt self.front_prompt = config.front_prompt self.sd = StableDiffusion(self.sd_device) - + if config.location_based_prompting: self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "") self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "") @@ -116,59 +116,59 @@ class DreamfusionPipeline(VanillaPipeline): metrics_dict = self.model.get_metrics_dict(model_outputs, batch) loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) - train_output = ( - model_outputs["train_output"] - .view(1, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3) - .permute(0, 3, 1, 2) - ) + # train_output = ( + # model_outputs["train_output"] + # .view(1, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3) + # .permute(0, 3, 1, 2) + # ) - if self.config.location_based_prompting: - if self.config.interpolated_prompting: - horiz = batch["central"].to(self.sd_device) - vert = max(batch["vertical"].to(self.sd_device), 0) + # if self.config.location_based_prompting: + # if self.config.interpolated_prompting: + # horiz = batch["central"].to(self.sd_device) + # vert = max(batch["vertical"].to(self.sd_device), 0) - if batch["central"] > 0 or batch["central"] <= 90: - text_embedding = ( - (horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding - ) / 90.0 - elif batch["central"] > 90 and batch["central"] <= 180: - text_embedding = ( - (horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding - ) / 90.0 - elif batch["central"] > 180 and batch["central"] <= 270: - text_embedding = ( - (horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding - ) / 90.0 - else: # batch["central"] > 270 and batch["central"] <= 360: - text_embedding = ( - (horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding - ) / 90.0 + # if batch["central"] > 0 or batch["central"] <= 90: + # text_embedding = ( + # (horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding + # ) / 90.0 + # elif batch["central"] > 90 and batch["central"] <= 180: + # text_embedding = ( + # (horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding + # ) / 90.0 + # elif batch["central"] > 180 and batch["central"] <= 270: + # text_embedding = ( + # (horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding + # ) / 90.0 + # else: # batch["central"] > 270 and batch["central"] <= 360: + # text_embedding = ( + # (horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding + # ) / 90.0 - text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0 + # text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0 - else: - if batch["vertical"] < 40: - text_embedding = self.top_text_embedding - elif batch["central"] > 315 or batch["central"] <= 45: - text_embedding = self.front_text_embedding - elif batch["central"] > 45 and batch["central"] <= 135: - text_embedding = self.side_text_embedding - elif batch["central"] > 135 and batch["central"] <= 225: - text_embedding = self.back_text_embedding - else: # batch["central"] > 225 and batch["central"] <= 315: - text_embedding = self.side_text_embedding - else: - text_embedding = self.base_text_embedding + # else: + # if batch["vertical"] < 40: + # text_embedding = self.top_text_embedding + # elif batch["central"] > 315 or batch["central"] <= 45: + # text_embedding = self.front_text_embedding + # elif batch["central"] > 45 and batch["central"] <= 135: + # text_embedding = self.side_text_embedding + # elif batch["central"] > 135 and batch["central"] <= 225: + # text_embedding = self.back_text_embedding + # else: # batch["central"] > 225 and batch["central"] <= 315: + # text_embedding = self.side_text_embedding + # else: + # text_embedding = self.base_text_embedding - with torch.autocast(device_type="cuda", dtype=torch.float16): - sds_loss = self.sd.sds_loss( - text_embedding.to(self.sd_device), - train_output.to(self.sd_device), - guidance_scale=int(self.config.guidance_scale), - grad_scaler=self.grad_scaler, - ) - - loss_dict["sds_loss"] = sds_loss.to(self.device) + # with torch.autocast(device_type="cuda", dtype=torch.float16): + # sds_loss = self.sd.sds_loss( + # text_embedding.to(self.sd_device), + # train_output.to(self.sd_device), + # guidance_scale=int(self.config.guidance_scale), + # grad_scaler=self.grad_scaler, + # ) + + # loss_dict["sds_loss"] = sds_loss.to(self.device) return model_outputs, loss_dict, metrics_dict From b89883562d8f73cb0a93e2db24b4890640a714e7 Mon Sep 17 00:00:00 2001 From: terrance Date: Thu, 16 Feb 2023 23:11:57 +0000 Subject: [PATCH 2/4] merging + remove dreamfusion pipeline file --- nerfstudio/configs/method_configs.py | 1 - nerfstudio/models/dreamfusion.py | 5 +- nerfstudio/pipelines/dreamfusion_pipeline.py | 184 ------------------- 3 files changed, 4 insertions(+), 186 deletions(-) delete mode 100644 nerfstudio/pipelines/dreamfusion_pipeline.py diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 270c7598e..e75f3a853 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -60,7 +60,6 @@ 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 from nerfstudio.plugins.registry import discover_methods diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index cea4cfb95..f96d16eb7 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -148,6 +148,8 @@ class DreamFusionModelConfig(ModelConfig): guidance_scale: float = 100 """guidance scale for sds loss""" stablediffusion_device: Optional[str] = None + """device for stable diffusion""" + sd_version: str = '1-5' class DreamFusionModel(Model): """DreamFusionModel Model @@ -165,6 +167,7 @@ class DreamFusionModel(Model): ) -> None: self.prompt = config.prompt self.cur_prompt = config.prompt + self.sd_version = config.sd_version self.initialize_density = config.initialize_density self.train_normals = False self.train_shaded = False @@ -192,7 +195,7 @@ class DreamFusionModel(Model): """Set the fields and modules""" super().populate_modules() - self.sd = StableDiffusion(self.sd_device, version=config.sd_version) + self.sd = StableDiffusion(self.sd_device, version=self.sd_version) self.text_embeddings = { "top_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", ""), "front_text_embedding": self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", ""), diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py deleted file mode 100644 index f70296b93..000000000 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ /dev/null @@ -1,184 +0,0 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Dreamfusion Pipeline and trainer""" - -from dataclasses import dataclass, field -from typing import Optional, Type - -import torch -from torch.cuda.amp.grad_scaler import GradScaler -from typing_extensions import Literal - -from nerfstudio.data.datamanagers.dreamfusion_datamanager import ( - DreamFusionDataManagerConfig, -) -from nerfstudio.generative.stable_diffusion import StableDiffusion -from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig - - -@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""" - # prompt: str = "A high-quality photo of a tree frog on a stump" - # """prompt for stable dreamfusion""" - - location_based_prompting: bool = True - """enables location based prompting""" - interpolated_prompting: bool = False - """enables interpolated location prompting""" - top_prompt: str = ", overhead view" - """appended to prompt for overhead view""" - side_prompt: str = ", side view" - """appended to prompt for side view""" - front_prompt: str = ", front view" - """appended to prompt for front view""" - back_prompt: str = ", back view" - """appended to prompt for back view""" - guidance_scale: float = 100 - """guidance scale for sds loss""" - sd_version: str = "1-5" - """Stable Diffusion version for dreamfusion""" - - stablediffusion_device: Optional[str] = None - - -class DreamfusionPipeline(VanillaPipeline): - """Dreamfusion pipeline""" - - config: DreamfusionPipelineConfig - - def __init__( - self, - config: DreamfusionPipelineConfig, - device: str, - test_mode: Literal["test", "val", "inference"] = "val", - world_size: int = 1, - local_rank: int = 0, - grad_scaler: Optional[GradScaler] = None, - ): - super().__init__(config, device, test_mode, world_size, local_rank, grad_scaler) - self.generative = True - self.cur_prompt = self.model.prompt - self.grad_scaler = grad_scaler - self.sd_device = ( - torch.device(device) - if self.config.stablediffusion_device is None - else torch.device(self.config.stablediffusion_device) - ) - self.location_based_prompting = config.location_based_prompting - self.top_prompt = config.top_prompt - self.side_prompt = config.side_prompt - self.back_prompt = config.back_prompt - self.front_prompt = config.front_prompt -<<<<<<< HEAD - self.sd = StableDiffusion(self.sd_device) - -======= - self.sd = StableDiffusion(self.sd_device, version=config.sd_version) - ->>>>>>> c81e24bd5ad0240bd94c9f1d98af55348cb2f27c - if config.location_based_prompting: - self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "") - self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "") - self.side_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", "") - self.back_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", "") - else: - self.base_text_embedding = self.sd.get_text_embeds(self.cur_prompt, "") - - def get_train_loss_dict(self, step: int): - - if self.world_size > 1 and step: - assert self.datamanager.train_sampler is not None - self.datamanager.train_sampler.set_epoch(step) - - if self.model.prompt != self.cur_prompt: - self.cur_prompt = self.model.prompt - if self.location_based_prompting: - self.top_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.top_prompt}", "") - self.front_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.front_prompt}", "") - self.side_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.side_prompt}", "") - self.back_text_embedding = self.sd.get_text_embeds(f"{self.cur_prompt}{self.back_prompt}", "") - else: - self.base_text_embedding = self.sd.get_text_embeds(self.cur_prompt, "") - - ray_bundle, batch = self.datamanager.next_train(step) - model_outputs = self.model(ray_bundle) - metrics_dict = self.model.get_metrics_dict(model_outputs, batch) - loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) - - # train_output = ( - # model_outputs["train_output"] - # .view(1, self.config.datamanager.train_resolution, self.config.datamanager.train_resolution, 3) - # .permute(0, 3, 1, 2) - # ) - - # if self.config.location_based_prompting: - # if self.config.interpolated_prompting: - # horiz = batch["central"].to(self.sd_device) - # vert = max(batch["vertical"].to(self.sd_device), 0) - - # if batch["central"] > 0 or batch["central"] <= 90: - # text_embedding = ( - # (horiz) * self.side_text_embedding + (90 - horiz) * self.front_text_embedding - # ) / 90.0 - # elif batch["central"] > 90 and batch["central"] <= 180: - # text_embedding = ( - # (horiz - 90) * self.back_text_embedding + (180 - horiz) * self.side_text_embedding - # ) / 90.0 - # elif batch["central"] > 180 and batch["central"] <= 270: - # text_embedding = ( - # (horiz - 180) * self.side_text_embedding + (270 - horiz) * self.back_text_embedding - # ) / 90.0 - # else: # batch["central"] > 270 and batch["central"] <= 360: - # text_embedding = ( - # (horiz - 270) * self.front_text_embedding + (360 - horiz) * self.side_text_embedding - # ) / 90.0 - - # text_embedding = (vert * text_embedding + (90 - vert) * self.top_text_embedding) / 90.0 - - # else: - # if batch["vertical"] < 40: - # text_embedding = self.top_text_embedding - # elif batch["central"] > 315 or batch["central"] <= 45: - # text_embedding = self.front_text_embedding - # elif batch["central"] > 45 and batch["central"] <= 135: - # text_embedding = self.side_text_embedding - # elif batch["central"] > 135 and batch["central"] <= 225: - # text_embedding = self.back_text_embedding - # else: # batch["central"] > 225 and batch["central"] <= 315: - # text_embedding = self.side_text_embedding - # else: - # text_embedding = self.base_text_embedding - - # with torch.autocast(device_type="cuda", dtype=torch.float16): - # sds_loss = self.sd.sds_loss( - # text_embedding.to(self.sd_device), - # train_output.to(self.sd_device), - # guidance_scale=int(self.config.guidance_scale), - # grad_scaler=self.grad_scaler, - # ) - - # loss_dict["sds_loss"] = sds_loss.to(self.device) - - return model_outputs, loss_dict, metrics_dict - - def forward(self): - """Not implemented since we only want the parameter saving of the nn module, but not forward()""" - raise NotImplementedError From 52ab31995306bbb7b25dc9725f3f9ebcb1c7695f Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 16 Feb 2023 15:45:54 -0800 Subject: [PATCH 3/4] Update configs --- nerfstudio/configs/method_configs.py | 8 ++++---- nerfstudio/model_components/ray_samplers.py | 8 ++++++-- nerfstudio/models/dreamfusion.py | 1 + 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index e75f3a853..dc9aba07f 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -352,15 +352,15 @@ method_configs["dreamfusion"] = TrainerConfig( ), model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, - distortion_loss_mult=10.0, + distortion_loss_mult=0.02, + interlevel_loss_mult=1.0, orientation_loss_mult=0.1, max_res=256, sphere_collider=True, initialize_density=False, random_background=True, - interlevel_loss_mult=1.0, - proposal_warmup=500, - proposal_update_every=5, + proposal_warmup=1000, + proposal_update_every=2, proposal_weights_anneal_max_num_iters=100, start_lambertian_training=1000, start_normals_training=500, diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 08ee1fd95..67201bb18 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -510,7 +510,7 @@ class VolumetricSampler(Sampler): class ProposalNetworkSampler(Sampler): """Sampler that uses a proposal network to generate samples. - + Args: num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step. num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model. @@ -518,6 +518,7 @@ class ProposalNetworkSampler(Sampler): single_jitter: Use a same random jitter for all samples along a ray. update_sched: A function that takes the iteration number of steps between updates. initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set. + pdf_histogram_padding: Padding to add to the histogram of the pdf. """ def __init__( @@ -528,6 +529,7 @@ class ProposalNetworkSampler(Sampler): single_jitter: bool = False, update_sched: Callable = lambda x: 1, initial_sampler: Optional[Sampler] = None, + pdf_histogram_padding: float = 0.01, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray @@ -542,7 +544,9 @@ class ProposalNetworkSampler(Sampler): self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) else: self.initial_sampler = initial_sampler - self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) + self.pdf_sampler = PDFSampler( + include_original=False, single_jitter=single_jitter, histogram_padding=pdf_histogram_padding + ) self._anneal = 1.0 self._steps_since_update = 0 diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index f96d16eb7..6eba9dc96 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -234,6 +234,7 @@ class DreamFusionModel(Model): single_jitter=self.config.use_single_jitter, update_sched=update_schedule, initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter), + pdf_histogram_padding=0.001, ) # renderers From a474aa3fb81d07d55beb27685b4e8d6cd5bf9514 Mon Sep 17 00:00:00 2001 From: terrance Date: Fri, 17 Feb 2023 00:16:25 +0000 Subject: [PATCH 4/4] add stable diff utils --- .../generative/stable_diffusion_utils.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 nerfstudio/generative/stable_diffusion_utils.py diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py new file mode 100644 index 000000000..82f61ff7a --- /dev/null +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -0,0 +1,61 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Stable diffusion utils +""" + +def get_text_embedding(batch, prompting_type, text_embeddings, sd_device): + """""" + + base = text_embeddings["base_text_embedding"] + front = text_embeddings["front_text_embedding"] + side = text_embeddings["side_text_embedding"] + top = text_embeddings["top_text_embedding"] + back = text_embeddings["back_text_embedding"] + + assert prompting_type in ("location_based", "interpolated", "base") + + if prompting_type == "location_based": + if batch["vertical"] < 40: + text_embedding = top + elif batch["central"] > 315 or batch["central"] <= 45: + text_embedding = front + elif batch["central"] > 45 and batch["central"] <= 135: + text_embedding = side + elif batch["central"] > 135 and batch["central"] <= 225: + text_embedding = back + else: # batch["central"] > 225 and batch["central"] <= 315: + text_embedding = side + + elif prompting_type == "interpolated": + horiz = batch["central"].to(sd_device) + vert = max(batch["vertical"].to(sd_device), 0) + + if batch["central"] > 0 or batch["central"] <= 90: + text_embedding = (horiz) * side + (90 - horiz) * front + elif batch["central"] > 90 and batch["central"] <= 180: + text_embedding = (horiz - 90) * back + (180 - horiz) * side + elif batch["central"] > 180 and batch["central"] <= 270: + text_embedding = (horiz - 180) * side + (270 - horiz) * back + else: # batch["central"] > 270 and batch["central"] <= 360: + text_embedding = (horiz - 270) * front + (360 - horiz) * side + + text_embedding = text_embedding / 90.0 + text_embedding = (vert * text_embedding + (90 - vert) * top) / 90.0 + + elif prompting_type == "base": + text_embedding = base + + return text_embedding