change the user in dockerfile

Dockerfile

@@ -3,8 +3,13 @@ FROM nvcr.io/nvidia/pytorch:23.08-py3
 LABEL maintainer="yuanze"
 LABEL email="[email protected]"
 
+RUN useradd -m -u 1000 user
+USER user
+
 WORKDIR /code
 
+RUN chown -R user:user /code
+
 # Install webp support
 RUN apt update && apt install libwebp-dev -y
 
@@ -24,6 +29,6 @@ COPY ./change_setup.txt /code/Pointnet2_PyTorch/pointnet2_ops_lib/setup.py
 RUN cd Pointnet2_PyTorch/pointnet2_ops_lib \
     && pip install .
 
-COPY . .
+COPY --chown=user:user . /code
 
 CMD ["python", "app.py"]

app.py

@@ -10,13 +10,14 @@ from feature_extractors.uni3d_embedding_encoder import Uni3dEmbeddingEncoder
 MAX_BATCH_SIZE = 16
 MAX_QUEUE_SIZE = 10
 MAX_K_RETRIEVAL = 20
+cache_dir = "./.cache"
 
-encoder = Uni3dEmbeddingEncoder()
+encoder = Uni3dEmbeddingEncoder(cache_dir)
 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")
-relation = load_dataset("VAST-AI/LD-T3D", split="full")
+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)
 
 @functools.lru_cache()
 def get_embedding(option, modality, angle=None):

feature_extractors/uni3d_embedding_encoder.py

@@ -279,17 +279,23 @@ def create_uni3d(uni3d_path):
     return model
 
 class Uni3dEmbeddingEncoder(FeatureExtractor):
-    def __init__(self, cache_dir=None, **kwargs) -> None:
+    def __init__(self, cache_dir, **kwargs) -> None:
         bpe_path = "utils/bpe_simple_vocab_16e6.txt.gz"
-        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")
+        # 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):
             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)