From 0628165fde040ff174e47d297283fe2072ac08ce Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Wed, 16 Nov 2022 10:21:21 -0800
Subject: [PATCH] Update nerfacc to 0.3.1

---
 nerfstudio/model_components/ray_samplers.py | 16 +++++++---------
 nerfstudio/models/instant_ngp.py            |  6 +++---
 pyproject.toml                              |  2 +-
 3 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py
index a65bd2128..159eb9f46 100644
--- a/nerfstudio/model_components/ray_samplers.py
+++ b/nerfstudio/model_components/ray_samplers.py
@@ -431,7 +431,7 @@ class VolumetricSampler(Sampler):
         near_plane: float = 0.0,
         far_plane: Optional[float] = None,
         cone_angle: float = 0.0,
-    ) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples", 2]]:
+    ) -> Tuple[RaySamples, TensorType["total_samples", 2]]:
         """Generate ray samples in a bounding box.
 
         Args:
@@ -442,9 +442,8 @@ class VolumetricSampler(Sampler):
             cone_angle: Cone angle for raymarching, set to 0 for uniform marching.
 
         Returns:
-            a tuple of (ray_samples, packed_info, ray_indices)
+            a tuple of (ray_samples, ray_indices)
             The ray_samples are packed, only storing the valid samples.
-            The packed_info contains all the information to recover packed samples into unpacked mode for rendering.
             The ray_indices contains the indices of the rays that each sample belongs to.
         """
 
@@ -455,7 +454,7 @@ class VolumetricSampler(Sampler):
         else:
             camera_indices = None
 
-        packed_info, starts, ends = nerfacc.ray_marching(
+        ray_indices, starts, ends = nerfacc.ray_marching(
             rays_o=rays_o,
             rays_d=rays_d,
             scene_aabb=self.scene_aabb,
@@ -470,13 +469,12 @@ class VolumetricSampler(Sampler):
         )
         num_samples = starts.shape[0]
         if num_samples == 0:
-            # create a single fake sample and update packed_info accordingly
-            # this says the last ray in packed_info has 1 sample, which starts and ends at 1
-            packed_info[-1, 1] = 1
+            # create a single fake sample
+            ray_indices[0] = 1
             starts = torch.ones((1, 1), dtype=starts.dtype, device=rays_o.device)
             ends = torch.ones((1, 1), dtype=ends.dtype, device=rays_o.device)
 
-        ray_indices = nerfacc.unpack_info(packed_info)
+        ray_indices = ray_indices.to(torch.long)
         origins = rays_o[ray_indices]
         dirs = rays_d[ray_indices]
         if camera_indices is not None:
@@ -493,7 +491,7 @@ class VolumetricSampler(Sampler):
             ),
             camera_indices=camera_indices,
         )
-        return ray_samples, packed_info, ray_indices
+        return ray_samples, ray_indices
 
 
 class ProposalNetworkSampler(Sampler):
diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py
index f30d2ad9a..6194a110a 100644
--- a/nerfstudio/models/instant_ngp.py
+++ b/nerfstudio/models/instant_ngp.py
@@ -166,7 +166,7 @@ class NGPModel(Model):
         num_rays = len(ray_bundle)
 
         with torch.no_grad():
-            ray_samples, packed_info, ray_indices = self.sampler(
+            ray_samples, ray_indices = self.sampler(
                 ray_bundle=ray_bundle,
                 near_plane=self.config.near_plane,
                 far_plane=self.config.far_plane,
@@ -178,7 +178,7 @@ class NGPModel(Model):
 
         # accumulation
         weights = nerfacc.render_weight_from_density(
-            packed_info=packed_info,
+            ray_indices=ray_indices,
             sigmas=field_outputs[FieldHeadNames.DENSITY],
             t_starts=ray_samples.frustums.starts,
             t_ends=ray_samples.frustums.ends,
@@ -201,7 +201,7 @@ class NGPModel(Model):
             "accumulation": accumulation,
             "depth": depth,
             "alive_ray_mask": alive_ray_mask,  # the rays we kept from sampler
-            "num_samples_per_ray": packed_info[:, 1],
+            "num_samples_per_ray": nerfacc.pack_info(ray_indices, num_rays),
         }
         return outputs
 
diff --git a/pyproject.toml b/pyproject.toml
index 52b5abfc5..36070754f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "mediapy==1.1.0",
     "msgpack==1.0.4",
     "msgpack_numpy==0.4.8",
-    "nerfacc==0.2.1",
+    "nerfacc==0.3.1",
     "open3d>=0.16.0",
     "opencv-python==4.6.0.66",
     "plotly==5.7.0",