Upload model

config.json

@@ -0,0 +1,35 @@
+{
+  "architectures": [
+    "CVLFaceRecognitionModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "wrapper.ModelConfig",
+    "AutoModel": "wrapper.CVLFaceRecognitionModel"
+  },
+  "conf": {
+    "color_space": "RGB",
+    "freeze": false,
+    "input_size": [
+      3,
+      112,
+      112
+    ],
+    "mask_ratio": 0.0,
+    "name": "base",
+    "output_dim": 512,
+    "rpe_config": {
+      "ctx_type": "rel_keypoint_splithead_unshared",
+      "method": "product",
+      "mode": "ctx",
+      "name": "KPRPE_shared",
+      "num_keypoints": 5,
+      "ratio": 1.9,
+      "rpe_on": "k",
+      "shared_head": true
+    },
+    "start_from": "",
+    "yaml_path": "models/vit_kprpe/configs/v1_base_kprpe_splithead_unshared.yaml"
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.33.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c6d37ea874c2f38ffc9a7f0e9247efc994c3fb5c12d044759ac294e19d127f7
+size 460344344

wrapper.py

@@ -0,0 +1,30 @@
+from transformers import PreTrainedModel
+from transformers import PretrainedConfig
+from omegaconf import OmegaConf
+from models import get_model
+import yaml
+
+class ModelConfig(PretrainedConfig):
+
+    def __init__(
+            self,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.conf = dict(yaml.safe_load(open('pretrained_model/model.yaml')))
+
+
+class CVLFaceRecognitionModel(PreTrainedModel):
+    config_class = ModelConfig
+
+    def __init__(self, cfg):
+        super().__init__(cfg)
+        model_conf = OmegaConf.create(cfg.conf)
+        self.model = get_model(model_conf)
+        self.model.load_state_dict_from_path('pretrained_model/model.pt')
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+
+