diff options
| author | AUTOMATIC <16777216c@gmail.com> | 2022-10-09 07:23:31 +0000 |
|---|---|---|
| committer | AUTOMATIC <16777216c@gmail.com> | 2022-10-09 07:23:31 +0000 |
| commit | c77c89cc83c618472ad352cf8a28fde28c3a1377 (patch) | |
| tree | a5b6e4a15eee659a05b5bdc32343178ce5f9bb9a /modules/sd_models.py | |
| parent | 050a6a798cec90ae2f881c2ddd3f0221e69907dc (diff) | |
| download | stable-diffusion-webui-gfx803-c77c89cc83c618472ad352cf8a28fde28c3a1377.tar.gz stable-diffusion-webui-gfx803-c77c89cc83c618472ad352cf8a28fde28c3a1377.tar.bz2 stable-diffusion-webui-gfx803-c77c89cc83c618472ad352cf8a28fde28c3a1377.zip | |
make main model loading and model merger use the same code
Diffstat (limited to 'modules/sd_models.py')
| -rw-r--r-- | modules/sd_models.py | 14 |
1 files changed, 9 insertions, 5 deletions
diff --git a/modules/sd_models.py b/modules/sd_models.py index cb3982b1..18fb8c2e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -122,6 +122,13 @@ def select_checkpoint(): return checkpoint_info
+def get_state_dict_from_checkpoint(pl_sd):
+ if "state_dict" in pl_sd:
+ return pl_sd["state_dict"]
+
+ return pl_sd
+
+
def load_model_weights(model, checkpoint_info):
checkpoint_file = checkpoint_info.filename
sd_model_hash = checkpoint_info.hash
@@ -131,11 +138,8 @@ def load_model_weights(model, checkpoint_info): pl_sd = torch.load(checkpoint_file, map_location="cpu")
if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}")
-
- if "state_dict" in pl_sd:
- sd = pl_sd["state_dict"]
- else:
- sd = pl_sd
+
+ sd = get_state_dict_from_checkpoint(pl_sd)
model.load_state_dict(sd, strict=False)
|