sygil-webui/scripts/webui/util/load_from_config.py
hlky 5a10bdbd47
proposed streamlit code organization changes
I want people of all skill levels to be able to contribute
This is one way the code could be split up with the aim of making it easy to understand and contribute especially for people on the lower end of the skill spectrum
All i've done is split things, I think renaming and reorganising is still needed
2022-09-08 17:57:38 +01:00

31 lines
785 B
Python

from util.imports import *
def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}")
sd = pl_sd["state_dict"]
model = instantiate_from_config(config.model)
m, u = model.load_state_dict(sd, strict=False)
if len(m) > 0 and verbose:
print("missing keys:")
print(m)
if len(u) > 0 and verbose:
print("unexpected keys:")
print(u)
model.cuda()
model.eval()
return model
def load_sd_from_config(ckpt, verbose=False):
print(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}")
sd = pl_sd["state_dict"]
return sd
#