mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 10:34:29 +08:00
undo config and merge in deepfloyd from main
This commit is contained in:
terrance
parent
cb10a77cae
commit
b1c1fa083b
27
.pre-commit-config.yaml
Normal file
27
.pre-commit-config.yaml
Normal file
@ -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
|
||||
@ -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)
|
||||
tyro.cli(generate_image)
|
||||
|
||||
@ -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
|
||||
|
||||
39
nerfstudio/generative/utils.py
Normal file
39
nerfstudio/generative/utils.py
Normal file
@ -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
|
||||
@ -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.
|
||||
]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user