From 9f491369a1f549c6aa569e1443035cf8e906f2d0 Mon Sep 17 00:00:00 2001
From: Bryce <github20210803@accounts.brycedrennan.com>
Date: Sat, 20 Jan 2024 07:57:47 -0800
Subject: [PATCH] fix: if weights are float32 but float16 was specified, still
 use float16

---
 imaginairy/utils/model_manager.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/imaginairy/utils/model_manager.py b/imaginairy/utils/model_manager.py
index bc9581a..dc0bc16 100644
--- a/imaginairy/utils/model_manager.py
+++ b/imaginairy/utils/model_manager.py
@@ -603,8 +603,8 @@ def load_sdxl_pipeline_from_diffusers_weights(
     text_encoder.load_state_dict(text_encoder_weights, assign=True)
     del text_encoder_weights
     lda = lda.to(device=device, dtype=torch.float32)
-    unet = unet.to(device=device)
-    text_encoder = text_encoder.to(device=device)
+    unet = unet.to(device=device, dtype=dtype)
+    text_encoder = text_encoder.to(device=device, dtype=dtype)
     if for_inpainting:
         StableDiffusionCls = StableDiffusion_XL_Inpainting
     else: