lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-04-30 03:01:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: handles when state_dict does not exist

Browse Source
This commit is contained in:
leko
2022-10-07 23:09:21 +08:00
committed by AUTOMATIC1111
parent 87db6f01cc
commit 616b7218f7
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
modules/sd_models.py
View File

@@ -122,7 +122,11 @@ def load_model_weights(model, checkpoint_file, sd_model_hash):
pl_sd = torch.load(checkpoint_file, map_location="cpu") pl_sd = torch.load(checkpoint_file, map_location="cpu")
if "global_step" in pl_sd: if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}") print(f"Global Step: {pl_sd['global_step']}")
sd = pl_sd["state_dict"]
if "state_dict" in pl_sd:
sd = pl_sd["state_dict"]
else:
sd = pl_sd
model.load_state_dict(sd, strict=False) model.load_state_dict(sd, strict=False)