| import os |
|
|
| import monai.networks.nets |
| import torch |
| from transformers import AutoConfig, AutoModel, PreTrainedModel |
| from vista3d_config import VISTA3DConfig |
|
|
|
|
| class VISTA3DModel(PreTrainedModel): |
| """VISTA3D model for hugging face""" |
|
|
| config_class = VISTA3DConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| if config.model_type == "VISTA3D": |
| self.network = monai.networks.nets.vista3d132( |
| encoder_embed_dim=config.encoder_embed_dim, |
| in_channels=config.input_channels, |
| ) |
|
|
| def forward(self, input): |
| return self.network(input) |
|
|
|
|
| def register_my_model(): |
| """Utility function to register VISTA3D model so that it can be instantiate by the AutoModel function.""" |
| AutoConfig.register("VISTA3D", VISTA3DConfig) |
| AutoModel.register(VISTA3DConfig, VISTA3DModel) |
|
|
|
|
| if __name__ == "__main__": |
| FILE_PATH = os.path.dirname(__file__) |
| MODEL_WEIGHT_PATH = os.path.join(FILE_PATH, "models/model.pt") |
| MODEL_PATH = os.path.join(FILE_PATH, "vista3d_pretrained_model") |
| config = VISTA3DConfig() |
| hugging_face_model = VISTA3DModel(config) |
| hugging_face_model.network.load_state_dict(torch.load(MODEL_WEIGHT_PATH)) |
| hugging_face_model.save_pretrained(MODEL_PATH) |
|
|
