--no-half-vae

1 files changed, 3 insertions, 0 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index e63d3c29..2cdcd84f 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -149,6 +149,7 @@ def load_model_weights(model, checkpoint_info):
         model.half()

 

     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

 

     vae_file = os.path.splitext(checkpoint_file)[0] + ".vae.pt"

     if os.path.exists(vae_file):

@@ -158,6 +159,8 @@ def load_model_weights(model, checkpoint_info):
 

         model.first_stage_model.load_state_dict(vae_dict)

 

+    model.first_stage_model.to(devices.dtype_vae)

+

     model.sd_model_hash = sd_model_hash

     model.sd_model_checkpoint = checkpoint_file

     model.sd_checkpoint_info = checkpoint_info