diff --git a/nerfstudio/data/datamanagers/lerf_datamanager.py b/nerfstudio/data/datamanagers/lerf_datamanager.py index c9419e82c..07ff1658e 100644 --- a/nerfstudio/data/datamanagers/lerf_datamanager.py +++ b/nerfstudio/data/datamanagers/lerf_datamanager.py @@ -194,7 +194,7 @@ class LERFDataManager(VanillaDataManager): # pylint: disable=abstract-method cfg={ "tile_size_range": (0.05, 0.5), "tile_size_res": 7, - "stride_scaler": 0.5, + "stride_scaler": 0.6, }, cache_dir=Path(f"outputs/{self.config.dataparser.data.name}"), ) diff --git a/nerfstudio/fields/lerf_field.py b/nerfstudio/fields/lerf_field.py index bd4d3cf02..2b9063c83 100644 --- a/nerfstudio/fields/lerf_field.py +++ b/nerfstudio/fields/lerf_field.py @@ -30,8 +30,8 @@ class LERFField(Field): self.clip_encs = torch.nn.ModuleList( [ # LERFField._get_encoding(16, 512, 16, indim=3) - LERFField._get_encoding(16, 128, 10, indim=3), - LERFField._get_encoding(128, 1024, 13, indim=3), + LERFField._get_encoding(16, 128, 12, indim=3), + LERFField._get_encoding(128, 512, 12, indim=3), ] ) tot_out_dims = sum([e.n_output_dims for e in self.clip_encs]) diff --git a/nerfstudio/models/lerf.py b/nerfstudio/models/lerf.py index 7dcd77646..7e8d15777 100644 --- a/nerfstudio/models/lerf.py +++ b/nerfstudio/models/lerf.py @@ -125,6 +125,35 @@ class LERFModel(NerfactoModel): return outputs + @torch.no_grad() + def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle) -> Dict[str, torch.Tensor]: + """Takes in camera parameters and computes the output of the model. + + LERF overrides this from base_model since we need to compute the max_across relevancy in multiple batches, + which are not independent since they need to use the same scale + Args: + camera_ray_bundle: ray bundle to calculate outputs over + """ + # TODO(justin) implement max across behavior + num_rays_per_chunk = self.config.eval_num_rays_per_chunk + image_height, image_width = camera_ray_bundle.origins.shape[:2] + num_rays = len(camera_ray_bundle) + outputs_lists = defaultdict(list) + for i in range(0, num_rays, num_rays_per_chunk): + start_idx = i + end_idx = i + num_rays_per_chunk + ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) + outputs = self.forward(ray_bundle=ray_bundle) + for output_name, output in outputs.items(): # type: ignore + outputs_lists[output_name].append(output) + outputs = {} + for output_name, outputs_list in outputs_lists.items(): + if not torch.is_tensor(outputs_list[0]): + # TODO: handle lists of tensors as well + continue + outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) # type: ignore + return outputs + def _get_outputs_nerfacto(self, ray_samples: RaySamples): field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals) weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) diff --git a/nerfstudio/pipelines/lerf_pipeline.py b/nerfstudio/pipelines/lerf_pipeline.py index 6f8fa7578..8faed1170 100644 --- a/nerfstudio/pipelines/lerf_pipeline.py +++ b/nerfstudio/pipelines/lerf_pipeline.py @@ -29,15 +29,19 @@ except ImportError: import torch + @dataclass class OpenCLIPNetworkConfig(cfg.InstantiateConfig): """Configuration for network instantiation""" _target: Type = field(default_factory=lambda: OpenCLIPNetwork) """target class to instantiate""" - clip_model_type: str = "ViT-B-16" - clip_model_pretrained: str = "laion2b_s34b_b88k" - clip_n_dims: int = 512 + # clip_model_type: str = "ViT-B-16" + # clip_model_pretrained: str = "laion2b_s34b_b88k" + # clip_n_dims: int = 512 + clip_model_type: str = "ViT-L-14" + clip_model_pretrained: str = "laion2b_s32b_b82k" + clip_n_dims: int = 768 negatives: Tuple[str] = ("object", "things", "stuff", "texture")