From b89883562d8f73cb0a93e2db24b4890640a714e7 Mon Sep 17 00:00:00 2001 From: terrance Date: Thu, 16 Feb 2023 23:11:57 +0000 Subject: [PATCH] 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