| from omegaconf import OmegaConf | |
| from ldm.util import instantiate_from_config | |
| import importlib | |
| import os | |
| import torch | |
| 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 | |
| def instantiate_from_config(config): | |
| if not "target" in config: | |
| if config == '__is_first_stage__': | |
| return None | |
| elif config == "__is_unconditional__": | |
| return None | |
| raise KeyError("Expected key `target` to instantiate.") | |
| return get_obj_from_str(config["target"])(**config.get("params", dict())) | |
| def get_obj_from_str(string, reload=False): | |
| module, cls = string.rsplit(".", 1) | |
| if reload: | |
| module_imp = importlib.import_module(module) | |
| importlib.reload(module_imp) | |
| return getattr(importlib.import_module(module, package=None), cls) | |
| 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 |