diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 5b222fa19..fcef56306 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -21,7 +21,6 @@ from dataclasses import dataclass from pathlib import Path from typing import List, Optional, Union, cast -import appdirs import mediapy import numpy as np import torch @@ -107,38 +106,13 @@ class StableDiffusion(nn.Module): sd_id = SD_IDENTIFIERS[version] pipe = StableDiffusionPipeline.from_pretrained(sd_id, torch_dtype=torch.float16) - assert pipe is not None + assert isinstance(pipe, StableDiffusionPipeline) pipe = pipe.to(self.device) 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" - if unet_traced_filename.exists(): - CONSOLE.print("Loading traced UNet.") - unet_traced = torch.jit.load(unet_traced_filename) - - class TracedUNet(torch.nn.Module): - """Jitted version of UNet""" - - def __init__(self): - super().__init__() - self.in_channels = pipe.unet.in_channels - self.device = pipe.unet.device - - def forward(self, latent_model_input, t, encoder_hidden_states): - """Forward pass""" - sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0] - return UNet2DConditionOutput(sample=sample) - - self.unet = TracedUNet() - del pipe.unet - else: - CONSOLE.print("[bold yellow] Warning: Loading UNet without JIT acceleration.") - CONSOLE.print(f"Run [yellow]ns-trace-sd --sd-version {version} [/yellow] for a speedup!") - self.unet = pipe.unet - self.unet.to(memory_format=torch.channels_last) + self.unet = pipe.unet + self.unet.to(memory_format=torch.channels_last) self.tokenizer = pipe.tokenizer self.text_encoder = pipe.text_encoder @@ -208,7 +182,7 @@ class StableDiffusion(nn.Module): noise = torch.randn_like(latents) latents_noisy = self.scheduler.add_noise(latents, noise, t) # type: ignore # pred noise - latent_model_input = torch.cat([latents_noisy] * 2) + latent_model_input = torch.cat((latents_noisy,) * 2) noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample # perform guidance diff --git a/nerfstudio/scripts/generative/__init__.py b/nerfstudio/scripts/generative/__init__.py deleted file mode 100644 index 21acafa24..000000000 --- a/nerfstudio/scripts/generative/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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. diff --git a/nerfstudio/scripts/generative/trace_stable_diffusion.py b/nerfstudio/scripts/generative/trace_stable_diffusion.py deleted file mode 100644 index 1110ae2ae..000000000 --- a/nerfstudio/scripts/generative/trace_stable_diffusion.py +++ /dev/null @@ -1,111 +0,0 @@ -# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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. - -"""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 nerfstudio.generative.stable_diffusion import SD_IDENTIFIERS -from nerfstudio.utils.rich_utils import CONSOLE - -try: - from diffusers import StableDiffusionPipeline -except ImportError: - CONSOLE.print("[bold red]Missing Stable Diffusion packages.") - CONSOLE.print(r"Install using [yellow]pip install nerfstudio\[gen][/yellow]") - CONSOLE.print(r"or [yellow]pip install -e .\[gen][/yellow] if installing from source.") - sys.exit(1) - - -@dataclass -class TraceSD: - """Traces the Stable Diffusion U-Net for better performance.""" - - sd_version: str - """Stable diffusion version.""" - save_dir: Path = Path(appdirs.user_data_dir("nerfstudio")) - """Directory to save the traced model.""" - - def main(self): - """Trace the Stable Diffusion U-Net and save it to disk.""" - # torch disable grad - torch.set_grad_enabled(False) - - 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) - - 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() - - # 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 - 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__": - entrypoint() diff --git a/pyproject.toml b/pyproject.toml index 2d487f02d..175807233 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,8 +66,8 @@ dependencies = [ # Generative related dependencies gen = [ - "diffusers==0.9.0", - "transformers==4.24.0", + "diffusers==0.16.1", + "transformers==4.29.2", ]