fix: if weights are float32 but float16 was specified, still use float16

4 months ago · e6a1c988c5
parent cf8a44b317
commit e6a1c988c5
1 changed files with 2 additions and 2 deletions
--- 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: