Spaces:
Runtime error
Runtime error
File size: 2,250 Bytes
cb33ff6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
""" Utiliy functions to load pre-trained models more easily """
import os
import pkg_resources
from omegaconf import OmegaConf
import torch
from huggingface_hub import hf_hub_download
from imagedream.ldm.util import instantiate_from_config
PRETRAINED_MODELS = {
"sd-v2.1-base-4view-ipmv": {
"config": "sd_v2_base_ipmv.yaml",
"repo_id": "Peng-Wang/ImageDream",
"filename": "sd-v2.1-base-4view-ipmv.pt",
},
"sd-v2.1-base-4view-ipmv-local": {
"config": "sd_v2_base_ipmv_local.yaml",
"repo_id": "Peng-Wang/ImageDream",
"filename": "sd-v2.1-base-4view-ipmv-local.pt",
},
}
def get_config_file(config_path):
cfg_file = pkg_resources.resource_filename(
"imagedream", os.path.join("configs", config_path)
)
if not os.path.exists(cfg_file):
raise RuntimeError(f"Config {config_path} not available!")
return cfg_file
def build_model(model_name, config_path=None, ckpt_path=None, cache_dir=None):
if (config_path is not None) and (ckpt_path is not None):
config = OmegaConf.load(config_path)
model = instantiate_from_config(config.model)
model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
return model
if not model_name in PRETRAINED_MODELS:
raise RuntimeError(
f"Model name {model_name} is not a pre-trained model. Available models are:\n- "
+ "\n- ".join(PRETRAINED_MODELS.keys())
)
model_info = PRETRAINED_MODELS[model_name]
# Instiantiate the model
print(f"Loading model from config: {model_info['config']}")
config_file = get_config_file(model_info["config"])
config = OmegaConf.load(config_file)
model = instantiate_from_config(config.model)
# Load pre-trained checkpoint from huggingface
if not ckpt_path:
ckpt_path = hf_hub_download(
repo_id=model_info["repo_id"],
filename=model_info["filename"],
cache_dir=cache_dir,
)
print(f"Loading model from cache file: {ckpt_path}")
model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
return model
|