in threestudio/models/guidance/deep_floyd_guidance.py [0:0]
def configure(self) -> None:
threestudio.info(f"Loading Deep Floyd ...")
self.weights_dtype = (
torch.float16 if self.cfg.half_precision_weights else torch.float32
)
# Create model
self.pipe = IFPipeline.from_pretrained(
self.cfg.pretrained_model_name_or_path,
text_encoder=None,
safety_checker=None,
watermarker=None,
feature_extractor=None,
requires_safety_checker=False,
variant="fp16" if self.cfg.half_precision_weights else None,
torch_dtype=self.weights_dtype,
cache_dir=self.cfg.cache_dir,
local_files_only=self.cfg.local_files_only
).to(self.device)
# Load lora weights
if self.cfg.lora_weights_path is not None:
self.pipe.load_lora_weights(self.cfg.lora_weights_path)
self.pipe.scheduler = self.pipe.scheduler.__class__.from_config(self.pipe.scheduler.config, variance_type="fixed_small")
if self.cfg.enable_memory_efficient_attention:
if parse_version(torch.__version__) >= parse_version("2"):
threestudio.info(
"PyTorch2.0 uses memory efficient attention by default."
)
elif not is_xformers_available():
threestudio.warn(
"xformers is not available, memory efficient attention is not enabled."
)
else:
threestudio.warn(
f"Use DeepFloyd with xformers may raise error, see https://github.com/deep-floyd/IF/issues/52 to track this problem."
)
self.pipe.enable_xformers_memory_efficient_attention()
if self.cfg.enable_sequential_cpu_offload:
self.pipe.enable_sequential_cpu_offload()
if self.cfg.enable_attention_slicing:
self.pipe.enable_attention_slicing(1)
if self.cfg.enable_channels_last_format:
self.pipe.unet.to(memory_format=torch.channels_last)
self.unet = self.pipe.unet.eval()
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.set_min_max_steps() # set to default value
if self.cfg.time_prior is not None:
m1, m2, s1, s2 = self.cfg.time_prior
weights = torch.cat(
(
torch.exp(
-((torch.arange(self.num_train_timesteps, m1, -1) - m1) ** 2)
/ (2 * s1**2)
),
torch.ones(m1 - m2 + 1),
torch.exp(
-((torch.arange(m2 - 1, 0, -1) - m2) ** 2) / (2 * s2**2)
),
)
)
weights = weights / torch.sum(weights)
self.time_prior_acc_weights = torch.cumsum(weights, dim=0)
self.alphas: Float[Tensor, "..."] = self.scheduler.alphas_cumprod.to(
self.device
)
self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(
self.device
)
self.grad_clip_val: Optional[float] = None
threestudio.info(f"Loaded Deep Floyd!")