From 8754beeb852dd1996974cc8c77a6a5892f5bcd0f Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 20 Feb 2023 10:05:09 -0800 Subject: [PATCH] cleanup --- nerfstudio/engine/trainer.py | 6 +- nerfstudio/generative/stable_diffusion.py | 13 ++- .../generative/stable_diffusion_utils.py | 3 +- nerfstudio/models/dreamfusion.py | 28 ++--- nerfstudio/pipelines/base_pipeline.py | 2 +- nerfstudio/viewer/server/viewer_utils.py | 2 +- pyproject.toml | 2 +- scripts/generative/trace_stable_diffusion.py | 105 +++++++++++------- scripts/trace_stablediff.py | 82 -------------- 9 files changed, 97 insertions(+), 146 deletions(-) delete mode 100644 scripts/trace_stablediff.py diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index b59ce5664..91722bf25 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -146,7 +146,11 @@ class Trainer: 'inference': does not load any dataset into memory """ self.pipeline = self.config.pipeline.setup( - device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank, grad_scaler=self.grad_scaler + device=self.device, + test_mode=test_mode, + world_size=self.world_size, + local_rank=self.local_rank, + grad_scaler=self.grad_scaler, ) self.optimizers = self.setup_optimizers() diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 0af0f23f6..a3955712b 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -49,11 +49,12 @@ logging.set_verbosity_error() IMG_DIM = 512 CONST_SCALE = 0.18215 SD_IDENTIFIERS = { - '1-5':'runwayml/stable-diffusion-v1-5', - '2-0':'stabilityai/stable-diffusion-2-base', - '2-1':'stabilityai/stable-diffusion-2-1-base' + "1-5": "runwayml/stable-diffusion-v1-5", + "2-0": "stabilityai/stable-diffusion-2-base", + "2-1": "stabilityai/stable-diffusion-2-1-base", } + @dataclass class UNet2DConditionOutput: """Class to hold traced model""" @@ -88,7 +89,7 @@ class StableDiffusion(nn.Module): num_train_timesteps: number of training timesteps """ - def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000, version='1-5') -> None: + def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000, version="1-5") -> None: super().__init__() self.device = device @@ -113,8 +114,8 @@ class StableDiffusion(nn.Module): pipe.enable_attention_slicing() # use jitted unet - filename_sd_id = sd_id.split('/')[-1] - unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / f'{filename_sd_id}_unet_traced.pt' + filename_sd_id = sd_id.split("/")[-1] + unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / f"{filename_sd_id}_unet_traced.pt" if unet_traced_filename.exists(): CONSOLE.print("Loading traced UNet.") unet_traced = torch.jit.load(unet_traced_filename) diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index 82f61ff7a..187b1a60a 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -16,6 +16,7 @@ Stable diffusion utils """ + def get_text_embedding(batch, prompting_type, text_embeddings, sd_device): """""" @@ -51,7 +52,7 @@ def get_text_embedding(batch, prompting_type, text_embeddings, sd_device): 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 diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 6eba9dc96..d3a658851 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -19,12 +19,12 @@ DreamFusion implementation. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Tuple, Type, Optional +from typing import Dict, List, Optional, Tuple, Type import numpy as np import torch -from torch.nn import Parameter from torch.cuda.amp.grad_scaler import GradScaler +from torch.nn import Parameter from nerfstudio.cameras.rays import RayBundle from nerfstudio.engine.callbacks import ( @@ -32,11 +32,11 @@ 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 +from nerfstudio.generative.stable_diffusion import StableDiffusion +from nerfstudio.generative.stable_diffusion_utils import get_text_embedding from nerfstudio.model_components.losses import ( MSELoss, distortion_loss, @@ -149,7 +149,8 @@ class DreamFusionModelConfig(ModelConfig): """guidance scale for sds loss""" stablediffusion_device: Optional[str] = None """device for stable diffusion""" - sd_version: str = '1-5' + sd_version: str = "1-5" + class DreamFusionModel(Model): """DreamFusionModel Model @@ -174,17 +175,17 @@ class DreamFusionModel(Model): 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.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.back_prompt = config.back_prompt self.front_prompt = config.front_prompt self.sd_device = ( - torch.device(kwargs['device']) + torch.device(kwargs["device"]) if config.stablediffusion_device is None else torch.device(config.stablediffusion_device) ) @@ -201,7 +202,7 @@ class DreamFusionModel(Model): "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, "") + "base_text_embedding": self.sd.get_text_embeds(self.cur_prompt, ""), } # setting up fields @@ -468,13 +469,14 @@ class DreamFusionModel(Model): "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, "") + "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) + 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) ) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index ac1b73dd6..bf4d400f8 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -237,7 +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 + grad_scaler=grad_scaler, ) self.model.to(device) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 2613934d6..bbdd1099b 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -494,7 +494,7 @@ class ViewerState: prompt = getattr(graph, "prompt", None) is_generative = prompt is not None - + if is_generative: self.vis["model/is_generative"].write(str(is_generative).lower()) diff --git a/pyproject.toml b/pyproject.toml index 99c6f8ca2..f1154fcd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,7 +100,7 @@ ns-install-cli = "scripts.completions.install:entrypoint" ns-process-data = "scripts.process_data:entrypoint" ns-download-data = "scripts.downloads.download_data:entrypoint" ns-train = "scripts.train:entrypoint" -ns-trace-sd = "scripts.trace_stablediff:entrypoint" +ns-trace-sd = "scripts.generative.trace_stable_diffusion:entrypoint" ns-viewer = "scripts.viewer.run_viewer:entrypoint" ns-eval = "scripts.eval:entrypoint" ns-render = "scripts.render:entrypoint" diff --git a/scripts/generative/trace_stable_diffusion.py b/scripts/generative/trace_stable_diffusion.py index a501bd424..eee069eb3 100644 --- a/scripts/generative/trace_stable_diffusion.py +++ b/scripts/generative/trace_stable_diffusion.py @@ -2,6 +2,7 @@ import functools import sys +from dataclasses import dataclass from pathlib import Path import appdirs @@ -9,6 +10,8 @@ import torch import tyro from rich.console import Console +from nerfstudio.generative.stable_diffusion import SD_IDENTIFIERS + CONSOLE = Console(width=120) try: @@ -20,55 +23,77 @@ except ImportError: sys.exit(1) -def jit_unet(save_dir: Path = Path(appdirs.user_data_dir("nerfstudio"))): - """Trace Stable Diffusion UNet for speed improvement +@dataclass +class TraceSD: + """Traces the Stable Diffusion U-Net for better performance.""" - Args: - save_dir: directory to save the traced model - """ + sd_version: str + """Stable diffusion version.""" + save_dir: Path = Path(appdirs.user_data_dir("nerfstudio")) + """Directory to save the traced model.""" - # torch disable grad - torch.set_grad_enabled(False) + def main(self): + """Trace the Stable Diffusion U-Net and save it to disk.""" + # torch disable grad + torch.set_grad_enabled(False) - # load inputs - def generate_inputs(): - sample = torch.randn(2, 4, 64, 64).half().cuda() - timestep = torch.rand(1).half().cuda() * 999 - encoder_hidden_states = torch.randn(2, 77, 768).half().cuda() - return sample, timestep, encoder_hidden_states + if self.sd_version not in SD_IDENTIFIERS: + CONSOLE.print(f"[bold red]Invalid Stable Diffusion version; choose from {set(SD_IDENTIFIERS.keys())}") + sys.exit(1) - pipe = StableDiffusionPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", - torch_dtype=torch.float16, - ).to("cuda") + sd_id = SD_IDENTIFIERS[self.sd_version] + pipe = StableDiffusionPipeline.from_pretrained( + sd_id, + torch_dtype=torch.float16, + ) + assert pipe is not None + pipe = pipe.to("cuda") + pipe.enable_attention_slicing() - pipe.enable_attention_slicing() + # load inputs + def generate_inputs(): + sample = torch.randn(2, 4, 64, 64).half().cuda() + timestep = torch.rand(1).half().cuda() * 999 - unet = pipe.unet - unet.eval() - unet.to(memory_format=torch.channels_last) # use channels_last memory format - unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default + # get shape of encoder hidden state (varies by SD version) + text_input = pipe.tokenizer("", return_tensors="pt") + _, w, d = pipe.text_encoder(text_input.input_ids.cuda())[0].shape - # warmup - inputs = None - for _ in range(3): - with torch.inference_mode(): - inputs = generate_inputs() - _ = unet(*inputs) + encoder_hidden_states = torch.randn(2, w, d).half().cuda() + return sample, timestep, encoder_hidden_states - # trace - print("tracing..") - assert inputs is not None - unet_traced = torch.jit.trace(unet, inputs) - unet_traced.eval() - print("done tracing") + unet = pipe.unet + unet.eval() + unet.to(memory_format=torch.channels_last) # use channels_last memory format + unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default - # save the traced model - unet_traced_filename = save_dir / "sd_unet_traced.pt" - if not unet_traced_filename.exists(): - unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) - unet_traced.save(unet_traced_filename) + # warmup + inputs = None + for _ in range(3): + with torch.inference_mode(): + inputs = generate_inputs() + _ = unet(*inputs) + + # trace + print("tracing..") + assert inputs is not None + unet_traced = torch.jit.trace(unet, inputs) + unet_traced.eval() # type: ignore + print("done tracing") + + # save the traced model + filename_sd_id = sd_id.split("/")[-1] + unet_traced_filename = self.save_dir / f"{filename_sd_id}_unet_traced.pt" + if not unet_traced_filename.exists(): + unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) + unet_traced.save(unet_traced_filename) # type: ignore + + +def entrypoint(): + """Entrypoint for use with pyproject scripts.""" + tyro.extras.set_accent_color("bright_yellow") + tyro.cli(TraceSD).main() if __name__ == "__main__": - tyro.cli(jit_unet) + entrypoint() diff --git a/scripts/trace_stablediff.py b/scripts/trace_stablediff.py deleted file mode 100644 index 44193307d..000000000 --- a/scripts/trace_stablediff.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Trace Stable Diffusion UNet for speed improvement""" - -import functools -import sys -from dataclasses import dataclass -from pathlib import Path - -import appdirs -import torch -import tyro -from diffusers import StableDiffusionPipeline -from rich.console import Console - -from nerfstudio.generative.stable_diffusion import SD_IDENTIFIERS - -CONSOLE = Console(width=120) - -@dataclass -class TraceSD: - """Traces the Stable Diffusion U-Net for better performance.""" - - sd_version: str - """Stable diffusion version.""" - - def main(self): - # torch disable grad - torch.set_grad_enabled(False) - - if self.sd_version not in SD_IDENTIFIERS.keys(): - CONSOLE.print(f"[bold red]Invalid Stable Diffusion version; choose from {set(SD_IDENTIFIERS.keys())}") - sys.exit(1) - - sd_id = SD_IDENTIFIERS[self.sd_version] - pipe = StableDiffusionPipeline.from_pretrained( - sd_id, - torch_dtype=torch.float16, - ).to("cuda") - pipe.enable_attention_slicing() - - # load inputs - def generate_inputs(): - sample = torch.randn(2, 4, 64, 64).half().cuda() - timestep = torch.rand(1).half().cuda() * 999 - - # get shape of encoder hidden state (varies by SD version) - text_input = pipe.tokenizer("", return_tensors="pt") - _, w, d = pipe.text_encoder(text_input.input_ids.cuda())[0].shape - - encoder_hidden_states = torch.randn(2, w, d).half().cuda() - return sample, timestep, encoder_hidden_states - - unet = pipe.unet - unet.eval() - unet.to(memory_format=torch.channels_last) # use channels_last memory format - unet.forward = functools.partial(unet.forward, return_dict=False) # set return_dict=False as default - - # warmup - for _ in range(3): - with torch.inference_mode(): - inputs = generate_inputs() - _ = unet(*inputs) - - # trace - print("tracing..") - unet_traced = torch.jit.trace(unet, inputs) - unet_traced.eval() - print("done tracing") - - # save the traced model - filename_sd_id = sd_id.split('/')[-1] - unet_traced_filename = Path(appdirs.user_data_dir("nerfstudio")) / f'{filename_sd_id}_unet_traced.pt' - if not unet_traced_filename.exists(): - unet_traced_filename.parent.mkdir(parents=True, exist_ok=True) - unet_traced.save(unet_traced_filename) - -def entrypoint(): - """Entrypoint for use with pyproject scripts.""" - tyro.extras.set_accent_color("bright_yellow") - tyro.cli(TraceSD).main() - -if __name__ == "__main__": - entrypoint()