diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 9b2359efd..16cec4e0f 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -476,19 +476,19 @@ method_configs["dreamfusion"] = TrainerConfig( "proposal_networks": { "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( - warmup_steps=4000, lr_final=5e-7, max_steps=20000, ramp="linear" + warmup_steps=4000, lr_final=5e-7, max_steps=30000, ramp="linear" ), }, "fields": { "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( - warmup_steps=4000, lr_final=5e-7, max_steps=20000, ramp="linear" + warmup_steps=4000, lr_final=5e-7, max_steps=30000, ramp="linear" ), }, "diffusion_model": { - "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=1e-4, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( - warmup_steps=4000, lr_final=5e-7, max_steps=20000, ramp="linear" + warmup_steps=4000, lr_final=5e-7, max_steps=30000, ramp="linear" ), }, }, diff --git a/nerfstudio/data/datamanagers/random_cameras_datamanager.py b/nerfstudio/data/datamanagers/random_cameras_datamanager.py index 4269b428a..368aa35f4 100644 --- a/nerfstudio/data/datamanagers/random_cameras_datamanager.py +++ b/nerfstudio/data/datamanagers/random_cameras_datamanager.py @@ -24,14 +24,15 @@ from typing import Dict, List, Tuple, Type, Union import torch from rich.progress import Console -from torch.nn import Parameter from torch import Tensor +from torch.nn import Parameter from typing_extensions import Literal from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.config_utils import to_immutable_dict -from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManagerConfig +from nerfstudio.data.datamanagers.base_datamanager import (DataManager, + DataManagerConfig) from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.scene_box import SceneBox from nerfstudio.data.utils.dataloaders import RandIndicesEvalDataloader @@ -72,7 +73,7 @@ def random_train_pose( focal_range: Tuple[float, float] = (0.75, 1.35), jitter_std: float = 0.01, center: Tuple[float, float, float] = (0, 0, 0), -) -> Tuple[Cameras, Tensor, Tensor]: +) -> Tuple[Cameras, Tensor, Tensor, Tensor]: """generate random poses from an orbit camera Args: size: batch size of generated poses. @@ -147,7 +148,7 @@ def random_train_pose( cy=resolution / 2, ).to(device) - return cameras, torch.rad2deg(vertical_rotation), torch.rad2deg(central_rotation) + return cameras, torch.rad2deg(vertical_rotation), torch.rad2deg(central_rotation), camera_to_worlds @dataclass @@ -218,7 +219,7 @@ class RandomCamerasDataManager(DataManager): # pylint: disable=abstract-method CONSOLE.print("[red] --data should not be used with the RandomCamerasDataManager[/red]") sys.exit(1) - cameras, _, _ = random_train_pose( + cameras, _, _, _ = random_train_pose( self.config.num_eval_angles, self.config.eval_resolution, device=self.device, @@ -249,7 +250,7 @@ class RandomCamerasDataManager(DataManager): # pylint: disable=abstract-method self.train_count += 1 horizontal_range = min((step / max(1, self.config.horizontal_rotation_warmup)), 1) * 180 - cameras, vertical_rotation, central_rotation = random_train_pose( + cameras, vertical_rotation, central_rotation, c2w = random_train_pose( self.config.train_images_per_batch, self.config.train_resolution, device=self.device, @@ -267,13 +268,14 @@ class RandomCamerasDataManager(DataManager): # pylint: disable=abstract-method "vertical": vertical_rotation, "central": central_rotation, "initialization": True, + "c2w": c2w, } def next_eval(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the eval dataloader.""" self.eval_count += 1 - cameras, vertical_rotation, central_rotation = random_train_pose( + cameras, vertical_rotation, central_rotation, c2w = random_train_pose( self.config.eval_images_per_batch, self.config.eval_resolution, device=self.device, @@ -288,7 +290,11 @@ class RandomCamerasDataManager(DataManager): # pylint: disable=abstract-method torch.tensor([[i] for i in range(self.config.train_images_per_batch)]) ).flatten() - return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation} + return ray_bundle, { + "vertical": vertical_rotation, + "central": central_rotation, + "c2w": c2w, + } def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]: for camera_ray_bundle, batch in self.eval_dataloader: diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index a66b9c1fd..47734ae12 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -16,7 +16,7 @@ # Modified from https://github.com/ashawkey/stable-dreamfusion/blob/main/nerf/sd.py -import copy +import os import sys from contextlib import contextmanager from dataclasses import dataclass @@ -30,9 +30,7 @@ import torch.nn.functional as F import tyro from jaxtyping import Float from torch import Tensor, nn -from torch.autograd import Variable from torch.cuda.amp.grad_scaler import GradScaler -from torchviz import make_dot from nerfstudio.generative.utils import _SDSGradient from nerfstudio.utils.rich_utils import CONSOLE @@ -106,11 +104,12 @@ class StableDiffusion(nn.Module): # num_train_timesteps=self.num_train_timesteps, # ) + os.environ["TOKENIZERS_PARALLELISM"] = "false" sd_id = SD_IDENTIFIERS[version] - pipe = StableDiffusionPipeline.from_pretrained(sd_id, torch_dtype=torch.float16).to(self.device) + pipe = StableDiffusionPipeline.from_pretrained(sd_id, torch_dtype=self.weights_dtype).to(self.device) assert isinstance(pipe, StableDiffusionPipeline) - pipe_lora = StableDiffusionPipeline.from_pretrained(sd_id, torch_dtype=torch.float16).to(self.device) + pipe_lora = StableDiffusionPipeline.from_pretrained(sd_id, torch_dtype=self.weights_dtype).to(self.device) del pipe_lora.vae pipe_lora.vae = pipe.vae @@ -158,7 +157,7 @@ class StableDiffusion(nn.Module): # p.requires_grad_(False) # # from https://github.com/threestudio-project/threestudio/commit/12d8f1c52f6ef4379db4e23261b274d36e12a531 - self.camera_embedding = TimestepEmbedding(16, 1280) + self.camera_embedding = TimestepEmbedding(12, 1280) self.unet_lora.class_embedding = self.camera_embedding lora_attn_procs = {} @@ -272,53 +271,6 @@ class StableDiffusion(nn.Module): return text_embeddings - def lora_loss( - self, - text_embeddings: Float[Tensor, "N max_length embed_dim"], - latents: Float[Tensor, "BS 4 64 64"], - guidance_scale: float = 100.0, - grad_scaler: Optional[GradScaler] = None, - ) -> torch.Tensor: - B = latents.shape[0] - t = torch.randint(0, int(self.scheduler.config.num_train_timesteps), [B], dtype=torch.long, device=self.device) - latents = latents.detach().repeat(1, 1, 1, 1) - - noise = torch.randn_like(latents) - latents_noisy = self.scheduler_lora.add_noise(latents, noise, t) # type: ignore - if self.scheduler_lora.config.prediction_type == "epsilon": - target = noise - elif self.scheduler_lora.config.prediction_type == "v_prediction": - target = self.scheduler_lora.get_velocity(latents, noise, t) - else: - raise ValueError(f"Unknown prediction type {self.scheduler_lora.config.prediction_type}") - - _, text_embeddings = text_embeddings.chunk(2) - # camera_condition = torch.zeros((B, 4, 4), device=self.device) - camera_condition = torch.eye(4, device=self.device)[None, ...] - - noise_pred = self.forward_unet( - self.unet_lora, - latents_noisy, - t, - encoder_hidden_states=text_embeddings, - class_labels=camera_condition.view(B, -1), - cross_attention_kwargs={"scale": 1.0}, - ) - - print("NOISE_PRED REQUIRES GRAD:", noise_pred.requires_grad, noise_pred.grad_fn) - - # q = [noise_pred.grad_fn] - # while len(q) > 0 and q[0] is not None: - # curr = q.pop(0) - # print(curr.name()) - # for next_fn in curr.next_functions: - # next_fn = next_fn[0] - # if next_fn is not None: - # q.append(next_fn) - # break - - return F.mse_loss(noise_pred.float(), target.float(), reduction="mean") - @torch.cuda.amp.autocast(enabled=False) def forward_unet( self, @@ -338,20 +290,49 @@ class StableDiffusion(nn.Module): cross_attention_kwargs=cross_attention_kwargs, ).sample.to(input_dtype) + def lora_loss( + self, + text_embeddings: Float[Tensor, "N max_length embed_dim"], + latents: Float[Tensor, "BS 4 64 64"], + camera_condition: Float[Tensor, "BS 3 4"], + guidance_scale: float = 100.0, + ) -> torch.Tensor: + B = latents.shape[0] + t = torch.randint(0, int(self.scheduler.config.num_train_timesteps), [B], dtype=torch.long, device=self.device) + latents = latents.detach().repeat(1, 1, 1, 1) + + noise = torch.randn_like(latents) + latents_noisy = self.scheduler_lora.add_noise(latents, noise, t) # type: ignore + if self.scheduler_lora.config.prediction_type == "epsilon": + target = noise + elif self.scheduler_lora.config.prediction_type == "v_prediction": + target = self.scheduler_lora.get_velocity(latents, noise, t) + else: + raise ValueError(f"Unknown prediction type {self.scheduler_lora.config.prediction_type}") + + _, text_embeddings = text_embeddings.chunk(2) + + noise_pred = self.forward_unet( + self.unet_lora, + latents_noisy, + t, + encoder_hidden_states=text_embeddings, + class_labels=camera_condition.view(B, -1), + cross_attention_kwargs={"scale": 1.0}, + ) + + return F.mse_loss(noise_pred.float(), target.float(), reduction="mean") + def vsd_loss( self, text_embeddings: Float[Tensor, "N max_length embed_dim"], latents: Float[Tensor, "BS 4 64 64"], + camera_condition: Float[Tensor, "BS 3 4"], guidance_scale: float = 1.0, grad_scaler: Optional[GradScaler] = None, ) -> torch.Tensor: B = latents.shape[0] - # tmp = self.lora_loss(text_embeddings, latents, guidance_scale=guidance_scale) - # print("LORA LOSS REQUIRES GRAD 1:", tmp.requires_grad) - - print("TEXT EMBEDDINGS SHAPE:", text_embeddings.shape) - with torch.no_grad(): # add noise t = torch.randint(self.min_step, self.max_step + 1, [B], dtype=torch.long, device=self.device) @@ -359,10 +340,8 @@ class StableDiffusion(nn.Module): latents_noisy = self.scheduler.add_noise(latents, noise, t) # type: ignore # pred noise latent_model_input = torch.cat([latents_noisy] * 2, dim=0) - # noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample with self.disable_unet_class_embedding(self.unet) as unet: cross_attention_kwargs = {"scale": 0.0} - # unet = copy.deepcopy(self.unet) noise_pred_pretrain = self.forward_unet( unet, latent_model_input, @@ -370,8 +349,6 @@ class StableDiffusion(nn.Module): encoder_hidden_states=text_embeddings, ) - camera_condition = torch.eye(4, device=self.device)[None, ...] - noise_pred_lora = self.forward_unet( self.unet_lora, latent_model_input, @@ -387,11 +364,7 @@ class StableDiffusion(nn.Module): cross_attention_kwargs={"scale": 1.0}, ) - # exit() - noise_pred_pretrain_uncond, noise_pred_pretrain_text = noise_pred_pretrain.chunk(2) - - # NOTE: guidance scale definition here is aligned with diffusers, but different from other guidance noise_pred_pretrain = noise_pred_pretrain_uncond + guidance_scale * ( noise_pred_pretrain_text - noise_pred_pretrain_uncond ) @@ -525,7 +498,7 @@ class StableDiffusion(nn.Module): Images """ - latents = 1 / CONST_SCALE * latents + latents = 1 / (self.vae.config.scaling_factor) * latents with torch.no_grad(): imgs = self.auto_encoder.decode(latents).sample @@ -546,7 +519,7 @@ class StableDiffusion(nn.Module): imgs = 2 * imgs - 1 posterior = self.vae.encode(imgs.to(self.weights_dtype)).latent_dist - latents = posterior.sample() * CONST_SCALE + latents = posterior.sample() * self.vae.config.scaling_factor return latents.to(input_dtype) @@ -607,6 +580,7 @@ class StableDiffusion(nn.Module): text_embeddings_vd: Float[Tensor, "N max_length embed_dim"], text_embeddings: Float[Tensor, "N max_length embed_dim"], image: Float[Tensor, "BS 3 H W"], + camera_condition: Float[Tensor, "BS 3 4"], guidance_scale: float = 1.0, grad_scaler: Optional[GradScaler] = None, ) -> torch.Tensor: @@ -615,61 +589,43 @@ class StableDiffusion(nn.Module): image = F.interpolate(image, (IMG_DIM, IMG_DIM), mode="bilinear") latents = self.imgs_to_latent(image) - loss_vsd = self.vsd_loss(text_embeddings_vd, latents, guidance_scale=guidance_scale) - loss_lora = self.lora_loss(text_embeddings, latents, guidance_scale=1.0) + loss_vsd = self.vsd_loss(text_embeddings_vd, latents, camera_condition, guidance_scale=guidance_scale) + loss_lora = self.lora_loss(text_embeddings, latents, camera_condition, guidance_scale=1.0) - print("LOSSES:", loss_vsd.item(), loss_lora.item()) - # exit() return loss_vsd + loss_lora def vsd_generate_image(prompt: str, negative: str = "", seed: int = 0, save_path: Path = Path("test_sd.png")): - # torch.manual_seed(seed) - # torch.cuda.manual_seed(seed) cuda_device = torch.device("cuda") sd = StableDiffusion(cuda_device, use_sds=False).to(cuda_device).float() im = torch.rand((1, 3, 512, 512)).float().to(cuda_device) im.requires_grad_(True) im.retain_grad() - # im = im.half() # .to(cuda_device) + text_embedding = sd.get_text_embeds(prompt, negative) text_embedding_vd = text_embedding - # WORKS - # text_embedding = torch.load("text_embeddings.pt").to(cuda_device) - # one, two = text_embedding.chunk(2) - # text_embedding = torch.cat([two, one], dim=0) - # text_embedding_vd = torch.load("text_embeddings_vd.pt").to(cuda_device) - # one, two = text_embedding_vd.chunk(2) - # text_embedding_vd = torch.cat([two, one], dim=0) + im_optimizer = torch.optim.Adam([im], lr=0.01, eps=1e-15) + sd_optimizer = torch.optim.Adam(list(sd.parameters()), lr=0.0001, eps=1e-15) - # AdamOptimizerConfig(lr=5e-3, eps=1e-15) - # optimizer = torch.optim.SGD([im] + list(sd.parameters()), lr=0.01, momentum=0.9, weight_decay=1e-4) - optimizer = torch.optim.Adam([im] + list(sd.parameters()), lr=0.01, eps=1e-15) + camera_condition = torch.eye(4, device=cuda_device)[None, ...] - # with torch.autocast(device_type="cuda", enabled=True): for i in range(5000): print(f"STEP: {i}; {im.std()}; {torch.cat([torch.flatten(p) for p in sd.unet_lora.parameters()]).std()}") - optimizer.zero_grad() + im_optimizer.zero_grad() + sd_optimizer.zero_grad() # loss = sd.sds_loss(text_embedding, im, guidance_scale=20.0) - loss = sd.vsd_forward(text_embedding_vd, text_embedding, im, guidance_scale=7.5) + loss = sd.vsd_forward(text_embedding_vd, text_embedding, im, camera_condition, guidance_scale=7.5) loss.backward() - # for name, param in sd.unet_lora.named_parameters(): - # if param.requires_grad and param.grad is not None: - # print(name, param.grad.norm()) - - optimizer.step() + im_optimizer.step() + sd_optimizer.step() if i % 10 == 0: mediapy.write_image(str(save_path), im.clip(0.0, 1.0).permute(0, 2, 3, 1).detach().cpu()[0]) - im = im.clip(0.0, 1.0).permute(0, 2, 3, 1).detach().cpu() - print("IM SHAPE:", im.shape) - mediapy.write_image(str(save_path), im[0]) - def generate_image( prompt: str, negative: str = "", seed: int = 0, steps: int = 50, save_path: Path = Path("test_sd.png") diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py index 3e2e793ae..3cfbbe19e 100644 --- a/nerfstudio/model_components/scene_colliders.py +++ b/nerfstudio/model_components/scene_colliders.py @@ -109,7 +109,6 @@ class AABBBoxCollider(SceneCollider): return ray_bundle -@torch_compile(dynamic=True, mode="reduce-overhead") def _intersect_with_sphere( rays_o: torch.Tensor, rays_d: torch.Tensor, center: torch.Tensor, radius: float = 1.0, near_plane: float = 0.0 ): diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 647015c25..a0d9cab41 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -339,11 +339,6 @@ class DreamFusionModel(Model): param_groups["proposal_networks"] = list(self.proposal_networks.parameters()) param_groups["fields"] = list(self.field.parameters()) # + list(self._diffusion_model.parameters()) param_groups["diffusion_model"] = list(self._diffusion_model.parameters()) - print("PARAMETERS:") - for name, param in self._diffusion_model.named_parameters(): - if param.requires_grad: - print(name) # , param.data) - # print(list(self._diffusion_model.parameters())) return param_groups def get_outputs(self, ray_bundle: RayBundle): # pylint: disable=too-many-statements @@ -478,6 +473,8 @@ class DreamFusionModel(Model): vertical_angle=batch["vertical"], horizontal_angle=batch["central"] ) + c2w = batch["c2w"] + train_output = ( outputs["train_output"] .view(1, int(outputs["train_output"].shape[0] ** 0.5), int(outputs["train_output"].shape[0] ** 0.5), 3) @@ -491,22 +488,16 @@ class DreamFusionModel(Model): # grad_scaler=self.grad_scaler, # ) - sds_loss = self._diffusion_model.vsd_loss( + sds_loss = self._diffusion_model.vsd_forward( + text_embedding.to(self.diffusion_device), text_embedding.to(self.diffusion_device), train_output.to(self.diffusion_device), - guidance_scale=int(self.guidance_scale), - grad_scaler=self.grad_scaler, - ) - - lora_loss = self._diffusion_model.lora_loss( - text_embedding.to(self.diffusion_device), - train_output.to(self.diffusion_device), - guidance_scale=int(1.0), + c2w.to(self.diffusion_device), + guidance_scale=int(7.5), grad_scaler=self.grad_scaler, ) loss_dict["sds_loss"] = sds_loss.to(self.device) - loss_dict["lora_loss"] = lora_loss.to(self.device) if self.training: loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss(