From f2a5cbe6f55592c4c5527b8e0bf99ea8d658f057 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 2 Nov 2022 14:41:29 +0300
Subject: [PATCH] fix #3986 breaking --no-half-vae

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

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 883639d..5075fad 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -183,11 +183,20 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
             model.to(memory_format=torch.channels_last)
 
         if not shared.cmd_opts.no_half:
+            vae = model.first_stage_model
+
+            # with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16
+            if shared.cmd_opts.no_half_vae:
+                model.first_stage_model = None
+
             model.half()
+            model.first_stage_model = vae
 
         devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16
         devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16
 
+        model.first_stage_model.to(devices.dtype_vae)
+
         if shared.opts.sd_checkpoint_cache > 0:
             # if PR #4035 were to get merged, restore base VAE first before caching
             checkpoints_loaded[checkpoint_key] = model.state_dict().copy()