diff options
| author | AUTOMATIC1111 <16777216c@gmail.com> | 2022-10-19 06:43:49 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-10-19 06:43:49 +0000 |
| commit | 05315d8a236e252221bbbdd9e8f459b8a31c3524 (patch) | |
| tree | 0bce187060568747888571fafedca4974fe17af3 /modules/sd_models.py | |
| parent | 9a33292ce41b01252cdb8ab6214a11d274e32fa0 (diff) | |
| parent | 1d4aa376e6111e90888a30ae24d2bcd7f978ec51 (diff) | |
| download | stable-diffusion-webui-gfx803-05315d8a236e252221bbbdd9e8f459b8a31c3524.tar.gz stable-diffusion-webui-gfx803-05315d8a236e252221bbbdd9e8f459b8a31c3524.tar.bz2 stable-diffusion-webui-gfx803-05315d8a236e252221bbbdd9e8f459b8a31c3524.zip | |
Merge branch 'master' into hot-reload-javascript
Diffstat (limited to 'modules/sd_models.py')
| -rw-r--r-- | modules/sd_models.py | 28 |
1 files changed, 25 insertions, 3 deletions
diff --git a/modules/sd_models.py b/modules/sd_models.py index 3aa21ec1..7ad6d474 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -122,11 +122,33 @@ def select_checkpoint(): return checkpoint_info
+chckpoint_dict_replacements = {
+ 'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.',
+ 'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.',
+ 'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.',
+}
+
+
+def transform_checkpoint_dict_key(k):
+ for text, replacement in chckpoint_dict_replacements.items():
+ if k.startswith(text):
+ k = replacement + k[len(text):]
+
+ return k
+
+
def get_state_dict_from_checkpoint(pl_sd):
if "state_dict" in pl_sd:
- return pl_sd["state_dict"]
+ pl_sd = pl_sd["state_dict"]
+
+ sd = {}
+ for k, v in pl_sd.items():
+ new_key = transform_checkpoint_dict_key(k)
+
+ if new_key is not None:
+ sd[new_key] = v
- return pl_sd
+ return sd
def load_model_weights(model, checkpoint_info):
@@ -141,7 +163,7 @@ def load_model_weights(model, checkpoint_info): print(f"Global Step: {pl_sd['global_step']}")
sd = get_state_dict_from_checkpoint(pl_sd)
- model.load_state_dict(sd, strict=False)
+ missing, extra = model.load_state_dict(sd, strict=False)
if shared.cmd_opts.opt_channelslast:
model.to(memory_format=torch.channels_last)
|