upgrade nerfacc to 0.3.x with CUB accel (#1187)

nerfstudio/model_components/ray_samplers.py

@@ -431,7 +431,7 @@ def forward(
         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",]]:
         """Generate ray samples in a bounding box.
 
         Args:
@@ -444,7 +444,6 @@ def forward(
         Returns:
             a tuple of (ray_samples, packed_info, 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 @@ def forward(
         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,
@@ -473,11 +472,10 @@ def forward(
         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
+            ray_indices = torch.zeros((1,), dtype=starts.dtype, device=rays_o.device)
             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)
         origins = rays_o[ray_indices]
         dirs = rays_d[ray_indices]
         if camera_indices is not None:
@@ -494,7 +492,7 @@ def forward(
             ),
             camera_indices=camera_indices,
         )
-        return ray_samples, packed_info, ray_indices
+        return ray_samples, ray_indices
 
 
 class ProposalNetworkSampler(Sampler):

nerfstudio/models/instant_ngp.py

@@ -170,7 +170,7 @@ def get_outputs(self, ray_bundle: RayBundle):
         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,
@@ -181,6 +181,7 @@ def get_outputs(self, ray_bundle: RayBundle):
         field_outputs = self.field(ray_samples)
 
         # accumulation
+        packed_info = nerfacc.pack_info(ray_indices, num_rays)
         weights = nerfacc.render_weight_from_density(
             packed_info=packed_info,
             sigmas=field_outputs[FieldHeadNames.DENSITY],

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.2",
     "open3d>=0.16.0",
     "opencv-python==4.6.0.66",
     "plotly==5.7.0",