tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-14 03:44:02 +08:00
Code Activity

cleanup

Browse Source
This commit is contained in:
Matt Tancik committed 2023-02-20 10:05:09 -08:00
1 parent 36be1fe4b0
commit 8754beeb85
9 files changed
+97 -146
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

nerfstudio/engine/trainer.py
+5 -1
View File
@@ -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()
nerfstudio/generative/stable_diffusion.py
+7 -6
View File
@@ -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)
nerfstudio/generative/stable_diffusion_utils.py
+2 -1
View File
@@ -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
nerfstudio/models/dreamfusion.py
+15 -13
View File
@@ -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)
)
nerfstudio/pipelines/base_pipeline.py
+1 -1
View File
@@ -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)
nerfstudio/viewer/server/viewer_utils.py
+1 -1
View File
@@ -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())
pyproject.toml
+1 -1
View File
@@ -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"
scripts/generative/trace_stable_diffusion.py
+65 -40
View File
@@ -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()
scripts/trace_stablediff.py
-82
View File
@@ -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()