(config, ckpt, verbose=True, freeze=True)
| 276 | |
| 277 | |
| 278 | def load_model_from_config(config, ckpt, verbose=True, freeze=True): |
| 279 | print(f"Loading model from {ckpt}") |
| 280 | if ckpt.endswith("ckpt"): |
| 281 | pl_sd = torch.load(ckpt, map_location="cpu") |
| 282 | if "global_step" in pl_sd: |
| 283 | print(f"Global Step: {pl_sd['global_step']}") |
| 284 | sd = pl_sd["state_dict"] |
| 285 | elif ckpt.endswith("safetensors"): |
| 286 | sd = load_safetensors(ckpt) |
| 287 | else: |
| 288 | raise NotImplementedError |
| 289 | |
| 290 | model = instantiate_from_config(config.model) |
| 291 | |
| 292 | m, u = model.load_state_dict(sd, strict=False) |
| 293 | |
| 294 | if len(m) > 0 and verbose: |
| 295 | print("missing keys:") |
| 296 | print(m) |
| 297 | if len(u) > 0 and verbose: |
| 298 | print("unexpected keys:") |
| 299 | print(u) |
| 300 | |
| 301 | if freeze: |
| 302 | for param in model.parameters(): |
| 303 | param.requires_grad = False |
| 304 | |
| 305 | model.eval() |
| 306 | return model |
| 307 | |
| 308 | |
| 309 | def get_configs_path() -> str: |
nothing calls this directly
no test coverage detected