tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Fix test

Browse Source
This commit is contained in:
Matt Tancik 2023-02-21 09:51:46 -08:00
parent 410b981205
commit 293ae8d643
1 changed files with 13 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
nerfstudio/generative/stable_diffusion.py
View File

@ -33,19 +33,16 @@ from torch.cuda.amp import custom_bwd, custom_fwd
from torch.cuda.amp.grad_scaler import GradScaler
from torchtyping import TensorType
CONSOLE = Console(width=120)
try:
from diffusers import PNDMScheduler, StableDiffusionPipeline
from transformers import logging
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)
SD_AVAILABLE = True
logging.set_verbosity_error()
except ImportError:
SD_AVAILABLE = False
CONSOLE = Console(width=120)
IMG_DIM = 512
CONST_SCALE = 0.18215
SD_IDENTIFIERS = {
@ -92,6 +89,14 @@ class StableDiffusion(nn.Module):
def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000, version="1-5") -> None:
super().__init__()
if not SD_AVAILABLE:
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)
logging.set_verbosity_error()
self.device = device
self.num_train_timesteps = num_train_timesteps