diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..4ccc2a61f --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,27 @@ +# See https://pre-commit.com for more information +# See https://pre-commit.com/hooks.html for more hooks +default_language_version: + python: python3 +repos: +- repo: local + hooks: + - id: add-license-headers + name: add-license-headers + entry: nerfstudio/scripts/licensing/license_headers.sh + language: script + files: '.*' + pass_filenames: false +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v3.2.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer +- repo: https://github.com/charliermarsh/ruff-pre-commit + # Ruff version. + rev: 'v0.0.267' + hooks: + - id: ruff +- repo: https://github.com/psf/black + rev: '23.3.0' + hooks: + - id: black diff --git a/nerfstudio/generative/deepfloyd.py b/nerfstudio/generative/deepfloyd.py index 3b5f3a41b..1401a58a1 100644 --- a/nerfstudio/generative/deepfloyd.py +++ b/nerfstudio/generative/deepfloyd.py @@ -1,48 +1,44 @@ -import sys -import gc -from dataclasses import dataclass -from pathlib import Path -from typing import List, Optional, Union +# 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. + +import gc +from pathlib import Path +from typing import List, Optional, Union, cast -import appdirs -import mediapy -import numpy as np import torch import torch.nn.functional as F import tyro -from rich.console import Console -from torch import nn -from torch.cuda.amp import custom_bwd, custom_fwd -from torch.cuda.amp.grad_scaler import GradScaler -from torchtyping import TensorType - from diffusers import IFPipeline -from transformers import T5EncoderModel, T5Tokenizer +from diffusers.pipelines.deepfloyd_if import IFPipelineOutput +from jaxtyping import Float +from PIL import Image +from torch import Generator, Tensor, nn +from torch.cuda.amp.grad_scaler import GradScaler + +from transformers import T5EncoderModel + +from nerfstudio.generative.utils import _SDSGradient IMG_DIM = 64 -class _SDSGradient(torch.autograd.Function): # pylint: disable=abstract-method - """Custom gradient function for SDS loss. Since it is already computed, we can just return it.""" - - @staticmethod - @custom_fwd - def forward(ctx, input_tensor, gt_grad): # pylint: disable=arguments-differ - del input_tensor - ctx.save_for_backward(gt_grad) - # Return magniture of gradient, not the actual loss. - return torch.mean(gt_grad**2) ** 0.5 - - @staticmethod - @custom_bwd - def backward(ctx, grad): # pylint: disable=arguments-differ - del grad - (gt_grad,) = ctx.saved_tensors - batch_size = len(gt_grad) - return gt_grad / batch_size, None class DeepFloyd(nn.Module): + """DeepFloyd diffusion model + Args: + device: device to use + """ - # load in model def __init__(self, device: Union[torch.device, str]): super().__init__() self.device = device @@ -64,27 +60,27 @@ class DeepFloyd(nn.Module): requires_safety_checker=False, variant="fp16", torch_dtype=torch.float16, - ).to(self.device) - + ) + assert isinstance(self.pipe, IFPipeline) + self.pipe = self.pipe.to(self.device) + self.pipe.enable_attention_slicing(1) - self.pipe.unet.to(memory_format=torch.channels_last) self.unet = self.pipe.unet + self.unet.to(memory_format=torch.channels_last) # type: ignore for p in self.unet.parameters(): p.requires_grad_(False) self.scheduler = self.pipe.scheduler - self.num_train_timesteps = self.scheduler.config.num_train_timesteps + self.num_train_timesteps = self.scheduler.config["num_train_timesteps"] self.min_step = int(self.num_train_timesteps * 0.02) self.max_step = int(self.num_train_timesteps * 0.98) + self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(self.device) - self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to( - self.device - ) - - def delete_text_encoder(self): + def delete_text_encoder(self): + """Delete text encoder from pipeline. T5 text encoder uses a lot of memory.""" del self.text_encoder del self.pipe gc.collect() @@ -99,12 +95,15 @@ class DeepFloyd(nn.Module): requires_safety_checker=False, variant="fp16", torch_dtype=torch.float16, - ).to(self.device) - + ) + assert isinstance(self.pipe, IFPipeline) + self.pipe = self.pipe.to(self.device) + self.pipe.enable_attention_slicing(1) - self.pipe.unet.to(memory_format=torch.channels_last) self.unet = self.pipe.unet + self.unet.to(memory_format=torch.channels_last) # type: ignore + for p in self.unet.parameters(): p.requires_grad_(False) @@ -112,20 +111,41 @@ class DeepFloyd(nn.Module): def get_text_embeds( self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]] - ) -> TensorType[2, "max_length", "embed_dim"]: + ) -> Float[Tensor, "2 max_length embed_dim"]: + """Get text embeddings for prompt and negative prompt + Args: + prompt: Prompt text + negative_prompt: Negative prompt text + Returns: + Text embeddings + """ prompt = [prompt] if isinstance(prompt, str) else prompt negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompt, negative_prompt=negative_prompt) + assert isinstance(self.pipe, IFPipeline) + with torch.no_grad(): + prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompt, negative_prompt=negative_prompt) + + assert isinstance(negative_embeds, Tensor) + assert isinstance(prompt_embeds, Tensor) return torch.cat([negative_embeds, prompt_embeds]) def sds_loss( self, - text_embeddings, - image, - guidance_scale, - grad_scaler, + text_embeddings: Float[Tensor, "N max_length embed_dim"], + image: Float[Tensor, "BS 3 H W"], + guidance_scale: float = 100.0, + grad_scaler: Optional[GradScaler] = None, ) -> torch.Tensor: + """Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/) + Args: + text_embeddings: Text embeddings + image: Rendered image + guidance_scale: How much to weigh the guidance + grad_scaler: Grad scaler + Returns: + The loss + """ with torch.autocast(device_type="cuda", enabled=False): image = F.interpolate(image.half(), (IMG_DIM, IMG_DIM), mode="bilinear", align_corners=False) t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) @@ -136,7 +156,7 @@ class DeepFloyd(nn.Module): noise = torch.randn_like(image) image_noisy = self.scheduler.add_noise(image, noise, t) # type: ignore # pred noise - image_model_input = torch.cat([image_noisy] * 2) + image_model_input = torch.cat((image_noisy,) * 2) noise_pred = self.unet(image_model_input, t, encoder_hidden_states=text_embeddings).sample # perform guidance @@ -153,8 +173,9 @@ class DeepFloyd(nn.Module): grad = torch.nan_to_num(grad) if grad_scaler is not None: - image = grad_scaler.scale(image) - loss = _SDSGradient.apply(image, grad) + loss = cast(Tensor, _SDSGradient.apply(grad_scaler.scale(image), grad)) + else: + loss = cast(Tensor, _SDSGradient.apply(image, grad)) return loss @@ -162,14 +183,15 @@ class DeepFloyd(nn.Module): self, prompts: Union[str, List[str]], negative_prompts: Union[str, List[str]] = "", - generator=None, + generator: Optional[Generator] = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, - ) -> np.ndarray: - """Generate an images from a prompts. + ) -> Image.Image: + """Generate an image from a prompt. Args: prompts: The prompt to generate an image from. negative_prompts: The negative prompt to generate an image from. + generator: Random seed. num_inference_steps: The number of inference steps to perform. guidance_scale: The scale of the guidance. latents: The latents to start from, defaults to random. @@ -179,12 +201,18 @@ class DeepFloyd(nn.Module): prompts = [prompts] if isinstance(prompts, str) else prompts negative_prompts = [negative_prompts] if isinstance(negative_prompts, str) else negative_prompts - # text_embeddings = self.get_text_embeds(prompts, negative_prompts) + + assert isinstance(self.pipe, IFPipeline) prompt_embeds, negative_embeds = self.pipe.encode_prompt(prompts, negative_prompt=negative_prompts) - image = self.pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator).images + model_output = self.pipe( + prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator + ) + assert isinstance(model_output, IFPipelineOutput) + image = model_output.images[0] return image + def generate_image( prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_deepfloyd.png") ): @@ -200,8 +228,9 @@ def generate_image( cuda_device = torch.device("cuda") with torch.no_grad(): df = DeepFloyd(cuda_device) - imgs = df.prompt_to_img(prompt, negative, generator, steps) - imgs[0].save(save_path) + img = df.prompt_to_img(prompt, negative, generator, steps) + img.save(save_path) + if __name__ == "__main__": - tyro.cli(generate_image) \ No newline at end of file + tyro.cli(generate_image) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index a3955712b..6202a9046 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -1,4 +1,4 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. +# 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. @@ -19,21 +19,21 @@ import sys from dataclasses import dataclass from pathlib import Path -from typing import List, Optional, Union +from typing import List, Optional, Union, cast -import appdirs import mediapy import numpy as np import torch import torch.nn.functional as F import tyro -from rich.console import Console -from torch import nn -from torch.cuda.amp import custom_bwd, custom_fwd +from jaxtyping import Float +from torch import Tensor, nn from torch.cuda.amp.grad_scaler import GradScaler -from torchtyping import TensorType -CONSOLE = Console(width=120) +from nerfstudio.utils.rich_utils import CONSOLE + +from nerfstudio.generative.utils import _SDSGradient + try: from diffusers import PNDMScheduler, StableDiffusionPipeline @@ -62,26 +62,6 @@ class UNet2DConditionOutput: sample: torch.FloatTensor -class _SDSGradient(torch.autograd.Function): # pylint: disable=abstract-method - """Custom gradient function for SDS loss. Since it is already computed, we can just return it.""" - - @staticmethod - @custom_fwd - def forward(ctx, input_tensor, gt_grad): # pylint: disable=arguments-differ - del input_tensor - ctx.save_for_backward(gt_grad) - # Return magniture of gradient, not the actual loss. - return torch.mean(gt_grad**2) ** 0.5 - - @staticmethod - @custom_bwd - def backward(ctx, grad): # pylint: disable=arguments-differ - del grad - (gt_grad,) = ctx.saved_tensors - batch_size = len(gt_grad) - return gt_grad / batch_size, None - - class StableDiffusion(nn.Module): """Stable Diffusion implementation Args: @@ -108,38 +88,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): # pylint: disable=no-self-use - """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 @@ -149,7 +104,7 @@ class StableDiffusion(nn.Module): def get_text_embeds( self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]] - ) -> TensorType[2, "max_length", "embed_dim"]: + ) -> Float[Tensor, "2 max_length embed_dim"]: """Get text embeddings for prompt and negative prompt Args: prompt: Prompt text @@ -185,8 +140,8 @@ class StableDiffusion(nn.Module): def sds_loss( self, - text_embeddings: TensorType["N", "max_length", "embed_dim"], - image: TensorType["BS", 3, "H", "W"], + text_embeddings: Float[Tensor, "N max_length embed_dim"], + image: Float[Tensor, "BS 3 H W"], guidance_scale: float = 100.0, grad_scaler: Optional[GradScaler] = None, ) -> torch.Tensor: @@ -209,7 +164,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 @@ -224,19 +179,19 @@ class StableDiffusion(nn.Module): if grad_scaler is not None: latents = grad_scaler.scale(latents) - loss = _SDSGradient.apply(latents, grad) + loss = cast(Tensor, _SDSGradient.apply(latents, grad)) return loss def produce_latents( self, - text_embeddings: TensorType["N", "max_length", "embed_dim"], + text_embeddings: Float[Tensor, "N max_length embed_dim"], height: int = IMG_DIM, width: int = IMG_DIM, num_inference_steps: int = 50, guidance_scale: float = 7.5, - latents: Optional[TensorType["BS", 4, "H", "W"]] = None, - ) -> TensorType["BS", 4, "H", "W"]: + latents: Optional[Float[Tensor, "BS 4 H W"]] = None, + ) -> Float[Tensor, "BS 4 H W"]: """Produce latents for a given text embedding Args: text_embeddings: Text embeddings @@ -251,7 +206,8 @@ class StableDiffusion(nn.Module): if latents is None: latents = torch.randn( - (text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device + (text_embeddings.shape[0] // 2, self.unet.config.in_channels, height // 8, width // 8), + device=self.device, ) self.scheduler.set_timesteps(num_inference_steps) # type: ignore @@ -276,7 +232,7 @@ class StableDiffusion(nn.Module): latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore return latents - def latents_to_img(self, latents: TensorType["BS", 4, "H", "W"]) -> TensorType["BS", 3, "H", "W"]: + def latents_to_img(self, latents: Float[Tensor, "BS 4 H W"]) -> Float[Tensor, "BS 3 H W"]: """Convert latents to images Args: latents: Latents to convert @@ -293,7 +249,7 @@ class StableDiffusion(nn.Module): return imgs - def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", 4, "H", "W"]: + def imgs_to_latent(self, imgs: Float[Tensor, "BS 3 H W"]) -> Float[Tensor, "BS 4 H W"]: """Convert images to latents Args: imgs: Images to convert diff --git a/nerfstudio/generative/utils.py b/nerfstudio/generative/utils.py new file mode 100644 index 000000000..e5996da6d --- /dev/null +++ b/nerfstudio/generative/utils.py @@ -0,0 +1,39 @@ +# 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. + + +import torch +from torch.cuda.amp import custom_bwd, custom_fwd + +"""Gradient helper functions for SDS loss""" + + +class _SDSGradient(torch.autograd.Function): + """Custom gradient function for SDS loss. Since it is already computed, we can just return it.""" + + @staticmethod + @custom_fwd + def forward(ctx, input_tensor, gt_grad): + del input_tensor + ctx.save_for_backward(gt_grad) + # Return magniture of gradient, not the actual loss. + return torch.mean(gt_grad**2) ** 0.5 + + @staticmethod + @custom_bwd + def backward(ctx, grad): + del grad + (gt_grad,) = ctx.saved_tensors + batch_size = len(gt_grad) + return gt_grad / batch_size, None diff --git a/pyproject.toml b/pyproject.toml index abe400912..c66c34b04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "nerfstudio" -version = "0.2.2" +version = "0.3.1" description = "All-in-one repository for state-of-the-art NeRFs" readme = "README.md" license = { text="Apache 2.0"} @@ -20,11 +20,11 @@ dependencies = [ "tyro>=0.3.31", "gdown>=4.6.0", "ninja>=1.10", - "functorch>=0.2.1", "h5py>=2.9.0", "imageio>=2.21.1", 'importlib-metadata>=6.0.0; python_version < "3.10"', "ipywidgets>=7.6", + "jaxtyping>=0.2.15", "jupyterlab>=3.3.4", "matplotlib>=3.5.3", "mediapy>=1.1.0", @@ -47,12 +47,11 @@ dependencies = [ "rich>=12.5.1", "scikit-image>=0.19.3", "tensorboard>=2.13.0", - "torch>=1.12.1,<2.0.0", + "torch>=1.13.1", + "torchvision>=0.14.1", "torchmetrics[image]>=0.9.3", - "torchtyping>=0.1.4", - "torchvision>=0.13.0", "typing_extensions>=4.4.0", - "viser>=0.0.5", + "viser>=0.0.10", "nuscenes-devkit>=1.1.1", "wandb>=0.13.3", "xatlas", @@ -67,18 +66,25 @@ dependencies = [ # Generative related dependencies gen = [ - "diffusers==0.9.0", - "transformers==4.24.0", + "diffusers==0.16.1", + "transformers==4.29.2", ] # Development packages dev = [ "black[jupyter]==23.3.0", - "pylint==2.13.4", + "pre-commit==3.3.2", "pytest==7.1.2", "pytest-xdist==2.5.0", "typeguard==2.13.3", + "ruff==0.0.267", + "sshconf==0.2.5", + "pycolmap==0.3.0", + "diffusers==0.16.1", + "opencv-stubs==0.0.7", + "transformers==4.29.2", + "pyright==1.1.308", ] # Documentation related packages @@ -114,7 +120,7 @@ ns-dev-sync-viser-message-defs = "nerfstudio.scripts.viewer.sync_viser_message_d [options] # equivalent to using --extra-index-url with pip, which is needed for specifying the CUDA version torch and torchvision dependency_links = [ - "https://download.pytorch.org/whl/cu113" + "https://download.pytorch.org/whl/cu118" ] [tool.setuptools.packages.find] @@ -123,38 +129,15 @@ include = ["nerfstudio*"] [tool.setuptools.package-data] "*" = ["*.cu", "*.json", "py.typed", "setup.bash", "setup.zsh"] -# black [tool.black] line-length = 120 -# pylint -[tool.pylint.messages_control] -max-line-length = 120 -generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"] -good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$" -jobs = 0 -ignored-classes = ["TensorDataclass"] - -disable = [ - "duplicate-code", - "fixme", - "logging-fstring-interpolation", - "too-many-arguments", - "too-many-branches", - "too-many-instance-attributes", - "too-many-locals", - "too-many-statements", - "unnecessary-ellipsis", -] - -#pytest [tool.pytest.ini_options] -addopts = "-n=4 --typeguard-packages=nerfstudio --torchtyping-patch-typeguard --disable-warnings" +addopts = "-n=4 --typeguard-packages=nerfstudio --disable-warnings" testpaths = [ "tests", ] -# pyright [tool.pyright] include = ["nerfstudio"] exclude = ["**/node_modules", @@ -162,10 +145,34 @@ exclude = ["**/node_modules", ] defineConstant = { DEBUG = true } -reportMissingImports = true +reportMissingImports = "warning" reportMissingTypeStubs = false reportPrivateImportUsage = false -reportUndefinedVariable = false pythonVersion = "3.8" pythonPlatform = "Linux" + +[tool.ruff] +select = [ + "E", # pycodestyle errors. + "F", # Pyflakes rules. + "PLC", # Pylint convention warnings. + "PLE", # Pylint errors. + "PLR", # Pylint refactor recommendations. + "PLW", # Pylint warnings. +] +ignore = [ + "E501", # Line too long. + "F722", # Forward annotation false positive from jaxtyping. Should be caught by pyright. + "F821", # Forward annotation false positive from jaxtyping. Should be caught by pyright. + "PLR2004", # Magic value used in comparison. + "PLR0915", # Too many statements. + "PLR0913", # Too many arguments. + "PLC0414", # Import alias does not rename variable. (this is used for exporting names) + "PLC1901", # Use falsey strings. + "PLR5501", # Use `elif` instead of `else if`. + "PLR0911", # Too many return statements. + "PLR0912", # Too many branches. + "PLW0603", # Globa statement updates are discouraged. + "PLW2901", # For loop variable overwritten. +]