minchul commited on
Commit
34a266d
1 Parent(s): 8e95cde

Upload model

Browse files
Files changed (3) hide show
  1. config.json +24 -0
  2. model.safetensors +3 -0
  3. wrapper.py +30 -0
config.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "CVLFaceRecognitionModel"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "wrapper.ModelConfig",
7
+ "AutoModel": "wrapper.CVLFaceRecognitionModel"
8
+ },
9
+ "conf": {
10
+ "color_space": "RGB",
11
+ "freeze": false,
12
+ "input_size": [
13
+ 3,
14
+ 112,
15
+ 112
16
+ ],
17
+ "name": "ir101",
18
+ "output_dim": 512,
19
+ "start_from": "",
20
+ "yaml_path": "models/iresnet_insightface/configs/v1_ir101.yaml"
21
+ },
22
+ "torch_dtype": "float32",
23
+ "transformers_version": "4.33.0"
24
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e90e0c8e24c94409fe53c2254a1343ef54f48a3a66b885b8ff69aefa21660ee7
3
+ size 260997808
wrapper.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PreTrainedModel
2
+ from transformers import PretrainedConfig
3
+ from omegaconf import OmegaConf
4
+ from models import get_model
5
+ import yaml
6
+
7
+ class ModelConfig(PretrainedConfig):
8
+
9
+ def __init__(
10
+ self,
11
+ **kwargs,
12
+ ):
13
+ super().__init__(**kwargs)
14
+ self.conf = dict(yaml.safe_load(open('pretrained_model/model.yaml')))
15
+
16
+
17
+ class CVLFaceRecognitionModel(PreTrainedModel):
18
+ config_class = ModelConfig
19
+
20
+ def __init__(self, cfg):
21
+ super().__init__(cfg)
22
+ model_conf = OmegaConf.create(cfg.conf)
23
+ self.model = get_model(model_conf)
24
+ self.model.load_state_dict_from_path('pretrained_model/model.pt')
25
+
26
+ def forward(self, x):
27
+ return self.model(x)
28
+
29
+
30
+