From 24f21583cdba2ae6cc51773b956c6ce068d3dfe4 Mon Sep 17 00:00:00 2001
From: AnyISalIn <anyisalin@gmail.com>
Date: Fri, 4 Aug 2023 11:43:27 +0800
Subject: [PATCH] fix: prevent cache model.state_dict() after model hijack

Signed-off-by: AnyISalIn <anyisalin@gmail.com>
---
 modules/sd_models.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 1d93d8935..ba15b4518 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -303,12 +303,13 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
         sd_models_xl.extend_sdxl(model)
 
     model.load_state_dict(state_dict, strict=False)
-    del state_dict
     timer.record("apply weights to model")
 
     if shared.opts.sd_checkpoint_cache > 0:
         # cache newly loaded model
-        checkpoints_loaded[checkpoint_info] = model.state_dict().copy()
+        checkpoints_loaded[checkpoint_info] = state_dict
+
+    del state_dict
 
     if shared.cmd_opts.opt_channelslast:
         model.to(memory_format=torch.channels_last)