Volinga model and NVOL documentation (#1486)

* added volinga model to NeRFStudio

* Removed unused dependencies

* Modified dockerfile to match main

* Corrected typo in core_code_checks.yml

* Corrected typo in core_code_checks.yml

* updated test with Volinga

* Added conversion to NVOL documentation

* [WIP] Refactor documentation and removing Volinga config

* Volinga use existing NerfactoModelConfig

* Added Unreal Engine documentation to extensions

* removed unused field class import

* Updated docs

* Documentation ready

* Update Dockerfile

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>

docs/extensions/unreal_engine.md

@@ -0,0 +1,40 @@
+# Exporting to Unreal Engine
+
+ ```{image} imgs/desolation_unreal.png
+ :width: 800 
+ :align: center 
+ :alt: NeRF in Unreal Engine 
+ ``` 
+
+## Overview
+
+NeRFStudio models can be used in Unreal Engine if they are converted to an NVOL file. NVOL is a new standard file format to store NeRFs in a fast and efficient way. NVOL files can be obtained from NeRFStudio checkpoints files (.ckpt) using the [Volinga Suite](https://volinga.ai/).
+
+:::{admonition} Note
+:class: warning
+Volinga Suite is not yet publicly available. If you are interested in converting your models to NVOL, sign up for the beta at [volinga.ai](https://volinga.ai).
+:::
+
+
+## Exporting your model to NVOL
+Currently NVOL file only supports Volinga model (which is based on nerfacto). You can train your model using the following command:
+
+```bash
+ns-train volinga --data /path/to/your/data --vis viewer
+```
+
+Once the training is done, you can find your checkpoint file in the `outputs/path-to-your-data/volinga` folder. Then, you can drag it to Volinga Suite to export it to NVOL.
+
+ ```{image} imgs/export_nvol.png 
+ :width: 400 
+ :align: center 
+ :alt: Nvol export in Voliga Suite 
+ ``` 
+
+Once the NVOL is ready, you can download it and use it in Unreal Engine.
+
+ ```{image} imgs/nvol_ready.png 
+ :width: 800 
+ :align: center 
+ :alt: NVOL ready to use 
+ ``` 

docs/index.md

@@ -58,8 +58,8 @@ Contributing<reference/contributing>
 ```{toctree}
 :hidden:
 :caption: Extensions
-
 extensions/blender_addon
+extensions/unreal_engine
 ```
 
 ```{toctree}

nerfstudio/configs/method_configs.py

@@ -63,6 +63,7 @@ method_configs: Dict[str, TrainerConfig] = {}
 descriptions = {
     "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.",
     "depth-nerfacto": "Nerfacto with depth supervision.",
+    "volinga": "Real-time rendering model from Volinga. Directly exportable to NVOL format at https://volinga.ai/",
     "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.",
     "instant-ngp-bounded": "Implementation of Instant-NGP. Recommended for bounded real and synthetic scenes",
     "mipnerf": "High quality model for bounded scenes. (slow)",
@@ -137,6 +138,47 @@ method_configs["depth-nerfacto"] = TrainerConfig(
     vis="viewer",
 )
 
+method_configs["volinga"] = TrainerConfig(
+    method_name="volinga",
+    steps_per_eval_batch=500,
+    steps_per_save=2000,
+    max_num_iterations=30000,
+    mixed_precision=True,
+    pipeline=VanillaPipelineConfig(
+        datamanager=VanillaDataManagerConfig(
+            dataparser=NerfstudioDataParserConfig(),
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=NerfactoModelConfig(
+            eval_num_rays_per_chunk=1 << 15,
+            hidden_dim=32,
+            hidden_dim_color=32,
+            hidden_dim_transient=32,
+            num_nerf_samples_per_ray=24,
+            proposal_net_args_list=[
+                {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": True},
+                {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": True},
+            ],
+        ),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
 method_configs["instant-ngp"] = TrainerConfig(
     method_name="instant-ngp",
     steps_per_eval_batch=500,

nerfstudio/fields/density_fields.py

@@ -21,7 +21,6 @@ from typing import Optional
 
 import numpy as np
 import torch
-from torch.nn.parameter import Parameter
 from torchtyping import TensorType
 
 from nerfstudio.cameras.rays import RaySamples
@@ -62,11 +61,15 @@ class HashMLPDensityField(Field):
         features_per_level=2,
     ) -> None:
         super().__init__()
-        self.aabb = Parameter(aabb, requires_grad=False)
+        self.register_buffer("aabb", aabb)
         self.spatial_distortion = spatial_distortion
         self.use_linear = use_linear
         growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
+        self.register_buffer("max_res", torch.tensor(max_res))
+        self.register_buffer("num_levels", torch.tensor(num_levels))
+        self.register_buffer("log2_hashmap_size", torch.tensor(log2_hashmap_size))
+
         config = {
             "encoding": {
                 "otype": "HashGrid",

nerfstudio/fields/nerfacto_field.py

@@ -116,9 +116,13 @@ class TCNNNerfactoField(Field):
     ) -> None:
         super().__init__()
 
-        self.aabb = Parameter(aabb, requires_grad=False)
+        self.register_buffer("aabb", aabb)
         self.geo_feat_dim = geo_feat_dim
 
+        self.register_buffer("max_res", torch.tensor(max_res))
+        self.register_buffer("num_levels", torch.tensor(num_levels))
+        self.register_buffer("log2_hashmap_size", torch.tensor(log2_hashmap_size))
+
         self.spatial_distortion = spatial_distortion
         self.num_images = num_images
         self.appearance_embedding_dim = appearance_embedding_dim

nerfstudio/models/nerfacto.py

@@ -73,6 +73,12 @@ class NerfactoModelConfig(ModelConfig):
     """How far along the ray to stop sampling."""
     background_color: Literal["random", "last_sample", "black", "white"] = "last_sample"
     """Whether to randomize the background color."""
+    hidden_dim: int = 64
+    """Dimension of hidden layers"""
+    hidden_dim_color: int = 64
+    """Dimension of hidden layers for color network"""
+    hidden_dim_transient: int = 64
+    """Dimension of hidden layers for transient network"""
     num_levels: int = 16
     """Number of levels of the hashmap for the base mlp."""
     max_res: int = 2048
@@ -93,8 +99,8 @@ class NerfactoModelConfig(ModelConfig):
     """Use the same proposal network. Otherwise use different ones."""
     proposal_net_args_list: List[Dict] = field(
         default_factory=lambda: [
-            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128},
-            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256},
+            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": False},
+            {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": False},
         ]
     )
     """Arguments for the proposal density fields."""
@@ -145,9 +151,12 @@ class NerfactoModel(Model):
         # Fields
         self.field = TCNNNerfactoField(
             self.scene_box.aabb,
+            hidden_dim=self.config.hidden_dim,
             num_levels=self.config.num_levels,
             max_res=self.config.max_res,
             log2_hashmap_size=self.config.log2_hashmap_size,
+            hidden_dim_color=self.config.hidden_dim_color,
+            hidden_dim_transient=self.config.hidden_dim_transient,
             spatial_distortion=scene_contraction,
             num_images=self.num_train_data,
             use_pred_normals=self.config.predict_normals,

tests/test_train.py

@@ -21,6 +21,7 @@ BLACKLIST = [
     "instant-ngp",
     "instant-ngp-bounded",
     "nerfacto",
+    "volinga",
     "phototourism",
     "depth-nerfacto",
     "nerfplayer-ngp",