From 39ffc38d97ab5ae97985f74afbdd8894c7ebea9d Mon Sep 17 00:00:00 2001
From: Chung Min Kim <chungmin99@gmail.com>
Date: Mon, 3 Apr 2023 18:36:49 +0000
Subject: [PATCH] Add openai-clip

---
 projects-lerf/lerf/encoders/clip_encoder.py |  4 +-
 projects-lerf/lerf/lerf_config.py           | 48 +++++++++++++++++++++
 projects-lerf/pyproject.toml                |  3 +-
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/projects-lerf/lerf/encoders/clip_encoder.py b/projects-lerf/lerf/encoders/clip_encoder.py
index eb238346a..ca8d46c06 100644
--- a/projects-lerf/lerf/encoders/clip_encoder.py
+++ b/projects-lerf/lerf/encoders/clip_encoder.py
@@ -14,7 +14,7 @@ from lerf.encoders.image_encoder import BaseImageEncoder, BaseImageEncoderConfig
 @dataclass
 class CLIPNetworkConfig(BaseImageEncoderConfig):
     _target: Type = field(default_factory=lambda: CLIPNetwork)
-    clip_model_type: str = "ViT-B-16"
+    clip_model_type: str = "ViT-B/16"
     clip_n_dims: int = 512
     negatives: Tuple[str] = ("object", "things", "stuff", "texture")
 
@@ -55,7 +55,7 @@ class CLIPNetwork(BaseImageEncoder):
 
     @property
     def name(self) -> str:
-        return "clip_{}".format(self.config.clip_model_type)
+        return "clip_openai_{}".format(self.config.clip_model_type)
 
     @property
     def embedding_dim(self) -> int:
diff --git a/projects-lerf/lerf/lerf_config.py b/projects-lerf/lerf/lerf_config.py
index 1f13d2d6e..cd3b2239b 100644
--- a/projects-lerf/lerf/lerf_config.py
+++ b/projects-lerf/lerf/lerf_config.py
@@ -17,6 +17,7 @@ from nerfstudio.plugins.types import MethodSpecification
 """
 Swap out the network config to use OpenCLIP or CLIP here.
 """
+from lerf.encoders.slip_encoder import SLIPNetworkConfig
 from lerf.encoders.clip_encoder import CLIPNetworkConfig
 from lerf.encoders.openclip_encoder import OpenCLIPNetworkConfig
 
@@ -159,3 +160,50 @@ lerf_method_lite = MethodSpecification(
     ),
     description="A lightweight version of LERF designed to work on smaller GPUs",
 )
+
+lerf_openai_method = MethodSpecification(
+    config=TrainerConfig(
+        method_name="lerf-openai",
+        steps_per_eval_batch=500,
+        steps_per_save=2000,
+        max_num_iterations=30000,
+        mixed_precision=True,
+        pipeline=LERFPipelineConfig(
+            datamanager=LERFDataManagerConfig(
+                dataparser=NerfstudioDataParserConfig(train_split_fraction=0.99),
+                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=LERFModelConfig(
+                eval_num_rays_per_chunk=1 << 15,
+                hashgrid_sizes=(19, 19),
+                hashgrid_layers=(12, 12),
+                hashgrid_resolutions=((16, 128), (128, 512)),
+                num_lerf_samples=24,
+            ),
+            network=CLIPNetworkConfig(
+                clip_model_type="ViT-B/16", clip_n_dims=512
+            )
+        ),
+        optimizers={
+            "proposal_networks": {
+                "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+                "scheduler": None,
+            },
+            "fields": {
+                "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15),
+                "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=30000),
+            },
+            "lerf": {
+                "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15),
+                "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=7000),
+            },
+        },
+        viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+        vis="viewer",
+    ),
+    description="Base config for CLIP-400M (openai)",
+)
\ No newline at end of file
diff --git a/projects-lerf/pyproject.toml b/projects-lerf/pyproject.toml
index f2d251593..e79a63a47 100644
--- a/projects-lerf/pyproject.toml
+++ b/projects-lerf/pyproject.toml
@@ -5,4 +5,5 @@ version = "0.1.0"
 [project.entry-points.'nerfstudio.method_configs']
 lerf = 'lerf.lerf_config:lerf_method'
 lerf_big = 'lerf.lerf_config:lerf_method_big'
-lerf_lite = 'lerf.lerf_config:lerf_method_lite'
\ No newline at end of file
+lerf_lite = 'lerf.lerf_config:lerf_method_lite'
+lerf_openai = 'lerf.lerf_config:lerf_openai_method'
\ No newline at end of file