remove local cache

app.py

@@ -10,14 +10,13 @@ from feature_extractors.uni3d_embedding_encoder import Uni3dEmbeddingEncoder
 MAX_BATCH_SIZE = 16
 MAX_QUEUE_SIZE = 10
 MAX_K_RETRIEVAL = 20
-cache_dir = "."
 
-encoder = Uni3dEmbeddingEncoder(cache_dir)
+encoder = Uni3dEmbeddingEncoder()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 source_id_list = torch.load("data/source_id_list.pt")
 source_to_id = {source_id: i for i, source_id in enumerate(source_id_list)}
-dataset = load_dataset("VAST-AI/LD-T3D", name=f"rendered_imgs_diag_above", split="base", cache_dir=cache_dir)
-relation = load_dataset("VAST-AI/LD-T3D", split="full", cache_dir=cache_dir)
+dataset = load_dataset("VAST-AI/LD-T3D", name=f"rendered_imgs_diag_above", split="base")
+relation = load_dataset("VAST-AI/LD-T3D", split="full")
 
 @functools.lru_cache()
 def get_embedding(option, modality, angle=None):

feature_extractors/uni3d_embedding_encoder.py

@@ -279,23 +279,17 @@ def create_uni3d(uni3d_path):
     return model
 
 class Uni3dEmbeddingEncoder(FeatureExtractor):
-    def __init__(self, cache_dir, **kwargs) -> None:
+    def __init__(self, cache_dir=None, **kwargs) -> None:
         bpe_path = "utils/bpe_simple_vocab_16e6.txt.gz"
-        # uni3d_path = os.path.join(cache_dir, "Uni3D", "modelzoo", "uni3d-g", "model.pt") # concat the subfolder as hf_hub_download will put it here
-        clip_path = os.path.join(cache_dir, "Uni3D", "open_clip_pytorch_model.bin")
-
-        # if not os.path.exists(uni3d_path):
-        #     hf_hub_download("BAAI/Uni3D", "model.pt", subfolder="modelzoo/uni3d-g", cache_dir=cache_dir, 
-        #                     local_dir=cache_dir + os.sep + "Uni3D")
-        if not os.path.exists(clip_path):
+        if cache_dir is None:
+            clip_path = hf_hub_download("timm/eva02_enormous_patch14_plus_clip_224.laion2b_s9b_b144k", "open_clip_pytorch_model.bin")
+        else:
+            clip_path = os.path.join(cache_dir, "Uni3D", "open_clip_pytorch_model.bin")
             hf_hub_download("timm/eva02_enormous_patch14_plus_clip_224.laion2b_s9b_b144k", "open_clip_pytorch_model.bin", 
-                            cache_dir=cache_dir, local_dir=cache_dir + os.sep + "Uni3D")
+                                        cache_dir=cache_dir, local_dir=cache_dir + os.sep + "Uni3D")
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.tokenizer = SimpleTokenizer(bpe_path)
-        # self.model = create_uni3d(uni3d_path)
-        # self.model.eval()
-        # self.model.to(self.device)
         self.clip_model, _, self.preprocess = open_clip.create_model_and_transforms(model_name="EVA02-E-14-plus", pretrained=clip_path)
         self.clip_model.to(self.device)