fix another bug found at #95

* simple check acceptable keys in the partial update routine
wkpark · Dec 28, 2023 · ef50c40 · ef50c40
1 parent 13e4b15
commit ef50c40
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scripts/model_mixer.py b/scripts/model_mixer.py
@@ -3681,7 +3681,8 @@ def fake_checkpoint(checkpoint_info, metadata, model_name, sha256, fake=True):
                 for k in weight_changed[prefix]:
                     # remove prefix, 'cond_stage_model.' or 'conditioner.' will be removed
                     key = k[len(prefix):]
-                    base_dict[key] = theta_0[k]
+                    if k in theta_0:
+                        base_dict[key] = theta_0[k]
                 if isxl:
                     shared.sd_model.conditioner.load_state_dict(base_dict, strict=False)
                 else: