| import os | |
| import torch | |
| from omegaconf import OmegaConf | |
| from ldm.util import instantiate_from_config | |
| def get_state_dict(d): | |
| return d.get('state_dict', d) | |
| def load_state_dict(ckpt_path, location='cpu'): | |
| _, extension = os.path.splitext(ckpt_path) | |
| if extension.lower() == ".safetensors": | |
| import safetensors.torch | |
| state_dict = safetensors.torch.load_file(ckpt_path, device=location) | |
| else: | |
| state_dict = get_state_dict(torch.load(ckpt_path, map_location=torch.device(location))) | |
| state_dict = get_state_dict(state_dict) | |
| print(f'Loaded state_dict from [{ckpt_path}]') | |
| return state_dict | |
| def create_model(config_path): | |
| config = OmegaConf.load(config_path) | |
| model = instantiate_from_config(config.model).cpu() | |
| print(f'Loaded model config from [{config_path}]') | |
| return model | |