diff --git a/.vscode/settings.json b/.vscode/settings.json deleted file mode 100644 index 5e3f61611..000000000 --- a/.vscode/settings.json +++ /dev/null @@ -1,125 +0,0 @@ -{ - "files.watcherExclude": { - "**/__pycache__/": true, - "**/.ipynb_checkpoints": true, - "**/*.ipynb": true, - "**/old*": true, - "**/runs*": true, - "**/temp*": true - }, - "editor.defaultFormatter": "esbenp.prettier-vscode", - "[javascript]": { - "editor.defaultFormatter": "esbenp.prettier-vscode" - }, - "[javascriptreact]": { - "editor.defaultFormatter": "esbenp.prettier-vscode" - }, - "[typescript]": { - "editor.defaultFormatter": "esbenp.prettier-vscode" - }, - "[typescriptreact]": { - "editor.defaultFormatter": "esbenp.prettier-vscode" - }, - "prettier.configPath": ".prettierrc.js", - "typescript.suggestionActions.enabled": false, - "javascript.suggestionActions.enabled": false, - "[python]": { - "editor.defaultFormatter": "ms-python.python", - "editor.codeActionsOnSave": { - "source.organizeImports": true - } - }, - "editor.formatOnSave": true, - "python.envFile": "${workspaceFolder}/.env", - "python.formatting.provider": "black", - "python.formatting.blackArgs": ["--line-length=120"], - "python.linting.pylintEnabled": true, - "python.linting.flake8Enabled": false, - "python.linting.enabled": true, - "python.sortImports.args": ["--src=${workspaceFolder}", "--profile", "black"], - "python.testing.unittestEnabled": false, - "python.testing.pytestEnabled": true, - "python.linting.ignorePatterns": [ - "**/site-packages/**/*.py", - ".vscode/*.py", - "external/**/*.py" - ], - "esbonio.sphinx.confDir": "", - // eslint extension options - "javascript.validate.enable": false, - "eslint.enable": true, - "eslint.validate": [ - "javascript", - "javascriptreact", - "typescript", - "typescriptreact" - ], - "C_Cpp.errorSquiggles": "Disabled", - "files.associations": { - "array": "cpp", - "bitset": "cpp", - "string_view": "cpp", - "initializer_list": "cpp", - "utility": "cpp", - "__hash_table": "cpp", - "__split_buffer": "cpp", - "deque": "cpp", - "iterator": "cpp", - "queue": "cpp", - "stack": "cpp", - "string": "cpp", - "unordered_map": "cpp", - "vector": "cpp", - "atomic": "cpp", - "*.tcc": "cpp", - "cctype": "cpp", - "chrono": "cpp", - "clocale": "cpp", - "cmath": "cpp", - "codecvt": "cpp", - "condition_variable": "cpp", - "cstdarg": "cpp", - "cstddef": "cpp", - "cstdint": "cpp", - "cstdio": "cpp", - "cstdlib": "cpp", - "cstring": "cpp", - "ctime": "cpp", - "cwchar": "cpp", - "cwctype": "cpp", - "exception": "cpp", - "algorithm": "cpp", - "filesystem": "cpp", - "functional": "cpp", - "memory": "cpp", - "memory_resource": "cpp", - "optional": "cpp", - "ratio": "cpp", - "system_error": "cpp", - "tuple": "cpp", - "type_traits": "cpp", - "fstream": "cpp", - "iomanip": "cpp", - "iosfwd": "cpp", - "iostream": "cpp", - "istream": "cpp", - "limits": "cpp", - "mutex": "cpp", - "new": "cpp", - "ostream": "cpp", - "sstream": "cpp", - "stdexcept": "cpp", - "streambuf": "cpp", - "thread": "cpp", - "typeinfo": "cpp", - "__nullptr": "cpp", - "__config": "cpp", - "__locale": "cpp", - "__bit_reference": "cpp", - "ios": "cpp", - "__atomic": "cpp", - "__node_handle": "cpp" - }, - "python.analysis.typeCheckingMode": "basic", - "python.analysis.diagnosticMode": "workspace" -} diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 3deaaf8fc..bdeabac03 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -50,6 +50,7 @@ from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig +from nerfstudio.models.dreamembedding import DreamEmbeddingModelConfig from nerfstudio.models.dreamfusion import DreamFusionModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel @@ -428,6 +429,53 @@ method_configs["dreamfusion"] = TrainerConfig( vis="viewer", ) +method_configs["dreamembed"] = TrainerConfig( + method_name="dreamembed", + experiment_name="dream-embed", + steps_per_eval_batch=50, + steps_per_eval_image=50, + steps_per_save=200, + max_num_iterations=20000, + mixed_precision=True, + pipeline=VanillaPipelineConfig( + generative=True, + datamanager=DreamFusionDataManagerConfig( + horizontal_rotation_warmup=4000, + ), + model=DreamEmbeddingModelConfig( + eval_num_rays_per_chunk=1 << 15, + distortion_loss_mult=0.02, + interlevel_loss_mult=100.0, + orientation_loss_mult=0.1, + max_res=256, + sphere_collider=True, + initialize_density=False, + random_background=True, + proposal_warmup=500, + proposal_update_every=0, + proposal_weights_anneal_max_num_iters=100, + start_normals_training=1000, + opacity_loss_mult=0.001, + positional_prompting="discrete", + guidance_scale=100, + ), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig( + warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear" + ), + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) + method_configs["nerfplayer-nerfacto"] = TrainerConfig( method_name="nerfplayer-nerfacto", pipeline=VanillaPipelineConfig( diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 34780a1eb..f3cb7964a 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -83,7 +83,7 @@ def random_train_pose( poses: [size, 4, 4] """ - vertical_rotation_range = [vertical_rotation_range[0] + 90, vertical_rotation_range[1] + 90] + vertical_rotation_range = (vertical_rotation_range[0] + 90, vertical_rotation_range[1] + 90) # This is the uniform sample on the part of the sphere we care about where 0 = 0 degrees and 1 = 360 degrees sampled_uniform = ( torch.rand(size) * (vertical_rotation_range[1] - vertical_rotation_range[0]) + vertical_rotation_range[0] @@ -157,7 +157,7 @@ class DreamFusionDataManagerConfig(DataManagerConfig): """Number of images per batch for training""" eval_images_per_batch: int = 1 """Number of images per batch for evaluation""" - radius_mean: float = 2.2 + radius_mean: float = 3.0 """Mean radius of camera orbit""" radius_std: float = 0.1 """Std of radius of camera orbit""" @@ -289,6 +289,23 @@ class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method ).flatten() return ray_bundle, {"vertical": vertical_rotation, "central": central_rotation} + + def next_eval_image(self, step:int) -> Tuple[int, RayBundle, Dict]: + cameras, vertical_rotation, central_rotation = random_train_pose( + self.config.eval_images_per_batch, + self.config.eval_resolution, + device=self.device, + radius_mean=self.config.radius_mean, + radius_std=self.config.radius_std, + focal_range=self.config.focal_range, + vertical_rotation_range=self.config.vertical_rotation_range, + jitter_std=self.config.jitter_std, + center=self.config.center, + ) + ray_bundle = cameras.generate_rays( + torch.tensor([[i] for i in range(self.config.train_images_per_batch)]) + ) + return 0, ray_bundle, {"vertical": vertical_rotation, "central": central_rotation} def get_train_rays_per_batch(self) -> int: return self.config.train_resolution**2 diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 752fd95a3..908f1a2a4 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -188,6 +188,42 @@ class StableDiffusion(nn.Module): return text_embeddings + def sds_loss_latent( + self, + text_embeddings: TensorType["N", "max_length", "embed_dim"], + latents: TensorType["BS", 3, "H", "W"], + guidance_scale: float = 100.0, + grad_scaler: Optional[GradScaler] = None, + ) -> torch.Tensor: + # Disable autocast when using diffusers pipeline. + with torch.autocast(device_type="cuda", enabled=False): + t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device) + latents = latents.half() + # predict the noise residual with unet, NO grad! + with torch.no_grad(): + # add noise + 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) + noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample + + # perform guidance + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_text + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # w(t), sigma_t^2 + w = 1 - self.alphas[t] + + grad = w * (noise_pred - noise) + grad = torch.nan_to_num(grad) + + if grad_scaler is not None: + latents = grad_scaler.scale(latents) + loss = _SDSGradient.apply(latents, grad) + + return loss + def sds_loss( self, text_embeddings: TensorType["N", "max_length", "embed_dim"], @@ -226,7 +262,7 @@ class StableDiffusion(nn.Module): # w(t), sigma_t^2 w = 1 - self.alphas[t] - grad = w * (noise_pred - noise) + grad = w * (noise_pred - noise) * 0.01 grad = torch.nan_to_num(grad) if grad_scaler is not None: diff --git a/nerfstudio/models/dreamembedding.py b/nerfstudio/models/dreamembedding.py index e97c30892..8a1871c1e 100644 --- a/nerfstudio/models/dreamembedding.py +++ b/nerfstudio/models/dreamembedding.py @@ -4,6 +4,7 @@ from __future__ import annotations from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple, Type +import mediapy import numpy as np import torch from torch.nn import Parameter @@ -32,15 +33,15 @@ from nerfstudio.model_components.ray_samplers import ( UniformSampler, ) from nerfstudio.model_components.renderers import FeatureRenderer - from nerfstudio.model_components.scene_colliders import AABBBoxCollider, SphereCollider from nerfstudio.model_components.shaders import LambertianShader, NormalsShader - from nerfstudio.models.base_model import Model, ModelConfig +from nerfstudio.models.dreamfusion import DreamFusionModel, DreamFusionModelConfig from nerfstudio.utils import colormaps, colors, math, misc + @dataclass -class DreamEmbeddingModelConfig(ModelConfig): +class DreamEmbeddingModelConfig(DreamFusionModelConfig): """DreamFusion model config""" _target: Type = field(default_factory=lambda: DreamEmbeddingModel) @@ -133,7 +134,7 @@ class DreamEmbeddingModelConfig(ModelConfig): sd_version: str = "1-5" -class DreamEmbeddingModel(Model): +class DreamEmbeddingModel(DreamFusionModel): """DreamEmbeddingModel Model Args: @@ -157,19 +158,15 @@ class DreamEmbeddingModel(Model): # setting up fields self.field = DreamEmbeddingField(self.scene_box.aabb, max_res=self.config.max_res) - self.feature_renderer = FeatureRenderer() - + self.renderer_feature = FeatureRenderer() def get_outputs(self, ray_bundle: RayBundle): # pylint: disable=too-many-statements # uniform sampling - background_rgb = self.field.get_background_feature(ray_bundle) + background_feature = self.field.get_background_feature(ray_bundle) ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) field_outputs = self.field(ray_samples, compute_normals=True) density = field_outputs[FieldHeadNames.DENSITY] - - # KEEP CHANIGNNG FROM HERE - if self.initialize_density: pos = ray_samples.frustums.get_positions() density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] @@ -181,21 +178,32 @@ class DreamEmbeddingModel(Model): accumulation = self.renderer_accumulation(weights) depth = self.renderer_depth(weights, ray_samples) - rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) + features = self.renderer_feature(features=field_outputs[FieldHeadNames.FEATURE], weights=weights) + # rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) accum_mask = torch.clamp((torch.nan_to_num(accumulation, nan=0.0)), min=0.0, max=1.0) accum_mask_inv = 1.0 - accum_mask - background = accum_mask_inv * background_rgb + background = accum_mask_inv * background_feature outputs = { - "rgb_only": rgb, - "background_rgb": background_rgb, + "feature_only": features, + "background_feature": background_feature, "background": background, "accumulation": accum_mask, "depth": depth, } + latents = accum_mask * features + background + outputs["latents"] = latents + + latents_input = latents.view(1, int(latents.shape[0] ** 0.5), int(latents.shape[0] ** 0.5), 4).half() + latents_input = latents_input.permute(0, 3, 1, 2) + rgb = self.sd.latents_to_img(latents_input).permute(0, 2, 3, 1) + rgb = rgb.view(512 * 512, 3) + rgb = rgb.to(torch.float32) + outputs["rgb"] = rgb + # These use a lot of GPU memory, so we avoid storing them for eval. if self.training: outputs["weights_list"] = weights_list @@ -205,55 +213,109 @@ class DreamEmbeddingModel(Model): outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) - pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) - - # lambertian shading - if self.config.random_light_source: # and self.training: - light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals) - else: - light_d = ray_bundle.origins[0] - light_d = math.safe_normalize(light_d) - - if (self.train_shaded and np.random.random_sample() > 0.75) or not self.training: - shading_weight = 0.9 - else: - shading_weight = 0.0 - - shaded, shaded_albedo = self.shader_lambertian( - rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False - ) - outputs["normals"] = self.shader_normals(normals, weights=accum_mask) - outputs["pred_normals"] = self.shader_normals(pred_normals, weights=accum_mask) - outputs["shaded"] = accum_mask * shaded - outputs["other_train_output"] = accum_mask * shaded_albedo + background - outputs["shaded_albedo"] = accum_mask * shaded_albedo - outputs["rgb"] = accum_mask * rgb + background - - if shading_weight > 0: - samp = np.random.random_sample() - if samp > 0.5 and not self.training: - outputs["train_output"] = outputs["shaded"] - elif samp < 0.2 and self.random_background: - rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) - outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv - else: - outputs["train_output"] = accum_mask * shaded_albedo + background - else: - outputs["train_output"] = outputs["rgb"] outputs["rendered_orientation_loss"] = orientation_loss( weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions ) - outputs["rendered_pred_normal_loss"] = pred_normal_loss( - weights.detach(), - field_outputs[FieldHeadNames.NORMALS].detach(), - field_outputs[FieldHeadNames.PRED_NORMALS], - ) - assert weights.shape[-1] == 1 if self.config.opacity_penalty: outputs["opacity_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult - return outputs \ No newline at end of file + return outputs + + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: + # Scaling metrics by coefficients to create the losses. + + loss_dict = {} + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) + if self.train_normals: + # orientation loss for computed normals + loss_dict["orientation_loss"] = self.config.orientation_loss_mult * torch.mean( + outputs["rendered_orientation_loss"] + ) + else: + loss_dict["orientation_loss"] = 0 + + if self.config.opacity_penalty: + loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean() + + if self.prompt != self.cur_prompt: + self.cur_prompt = self.prompt + self.text_embeddings.update_prompt( + base_prompt=self.cur_prompt, + top_prompt=self.cur_prompt + self.top_prompt, + side_prompt=self.cur_prompt + self.side_prompt, + back_prompt=self.cur_prompt + self.back_prompt, + front_prompt=self.cur_prompt + self.front_prompt, + ) + + text_embedding = self.text_embeddings.get_text_embedding( + vertical_angle=batch["vertical"], horizontal_angle=batch["central"] + ) + + latents = ( + outputs["latents"] + .view(1, int(outputs["latents"].shape[0] ** 0.5), int(outputs["latents"].shape[0] ** 0.5), 4) + .permute(0, 3, 1, 2) + ) + + # train_output = ( + # 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) + # ) + + # sds_loss = self.sd.sds_loss( + # text_embedding.to(self.sd_device), + # train_output.to(self.sd_device), + # guidance_scale=int(self.guidance_scale), + # grad_scaler=self.grad_scaler, + # ) + + sds_loss = self.sd.sds_loss_latent( + text_embedding.to(self.sd_device), + latents.to(self.sd_device), + guidance_scale=int(self.guidance_scale), + grad_scaler=self.grad_scaler, + ) + + loss_dict["sds_loss"] = sds_loss.to(self.device) + + if self.training: + loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + return loss_dict + + def get_image_metrics_and_images( + self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] + ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: + acc = colormaps.apply_colormap(outputs["accumulation"]) + depth = colormaps.apply_depth_colormap( + outputs["depth"], + accumulation=outputs["accumulation"], + ) + prop_depth_0 = colormaps.apply_depth_colormap( + outputs["prop_depth_0"], + accumulation=outputs["accumulation"], + ) + prop_depth_1 = colormaps.apply_depth_colormap( + outputs["prop_depth_1"], + accumulation=outputs["accumulation"], + ) + + metrics_dict = {} + images_dict = { + "img": outputs["rgb"].reshape(512, 512, 3), # this is a bit weird + "accumulation": acc, + "depth": depth, + "prop_depth_0": prop_depth_0, + "prop_depth_1": prop_depth_1, + "normals": outputs["normals"], + } + return metrics_dict, images_dict \ No newline at end of file diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 720d1d449..83d06e371 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -283,7 +283,7 @@ class WandbWriter(Writer): """WandDB Writer Class""" def __init__(self, log_dir: Path): - wandb.init(project="nerfstudio-project", dir=str(log_dir), reinit=True) + wandb.init(project="dream-embed", dir=str(log_dir), reinit=True) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: image = torch.permute(image, (2, 0, 1))