| from transformers import PreTrainedModel |
| from transformers import PretrainedConfig |
| from omegaconf import OmegaConf |
| from aligners import get_aligner |
| import yaml |
|
|
| class ModelConfig(PretrainedConfig): |
|
|
| def __init__( |
| self, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.conf = dict(yaml.safe_load(open('pretrained_model/model.yaml'))) |
|
|
|
|
| class CVLFaceAlignmentModel(PreTrainedModel): |
| config_class = ModelConfig |
|
|
| def __init__(self, cfg): |
| super().__init__(cfg) |
| model_conf = OmegaConf.create(cfg.conf) |
| self.model = get_aligner(model_conf) |
| self.model.load_state_dict_from_path('pretrained_model/model.pt') |
|
|
| def forward(self, *args, **kwargs): |
| return self.model(*args, **kwargs) |
|
|
|
|
|
|
|
|
