| 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 |
