(model_config=None, model_ckpt_path=None, pretrained_name=None, model_half=False)
| 15 | |
| 16 | |
| 17 | def load_model(model_config=None, model_ckpt_path=None, pretrained_name=None, model_half=False): |
| 18 | if pretrained_name is not None: |
| 19 | print(f"Loading pretrained model {pretrained_name}") |
| 20 | model, model_config = get_pretrained_model(pretrained_name) |
| 21 | |
| 22 | elif model_config is not None and model_ckpt_path is not None: |
| 23 | print(f"Creating model from config") |
| 24 | model = create_model_from_config(model_config) |
| 25 | |
| 26 | print(f"Loading model checkpoint from {model_ckpt_path}") |
| 27 | copy_state_dict(model, load_ckpt_state_dict(model_ckpt_path)) |
| 28 | |
| 29 | model.eval().requires_grad_(False) |
| 30 | |
| 31 | if model_half: |
| 32 | model.to(torch.float16) |
| 33 | |
| 34 | print("Done loading model") |
| 35 | |
| 36 | return model, model_config |
| 37 | |
| 38 | |
| 39 | class PreEncodedLatentsInferenceWrapper(pl.LightningModule): |
no test coverage detected