From c77c89cc83c618472ad352cf8a28fde28c3a1377 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 9 Oct 2022 10:23:31 +0300
Subject: make main model loading and model merger use the same code

---
 modules/sd_models.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

(limited to 'modules/sd_models.py')

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)
 
-- 
cgit v1.2.1