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/extras.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

(limited to 'modules/extras.py')

diff --git a/modules/extras.py b/modules/extras.py
index 1d9e64e5..ef6e6de7 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -169,9 +169,9 @@ def run_modelmerger(primary_model_name, secondary_model_name, interp_method, int
 
     print(f"Loading {secondary_model_info.filename}...")
     secondary_model = torch.load(secondary_model_info.filename, map_location='cpu')
-   
-    theta_0 = primary_model['state_dict']
-    theta_1 = secondary_model['state_dict']
+
+    theta_0 = sd_models.get_state_dict_from_checkpoint(primary_model)
+    theta_1 = sd_models.get_state_dict_from_checkpoint(secondary_model)
 
     theta_funcs = {
         "Weighted Sum": weighted_sum,
-- 
cgit v1.2.3