From 11e9166d0b00fe3b1e6ebb0a3d1db50ce7a56d58 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 10 Apr 2025 18:04:17 +0300
Subject: [PATCH] Fix low vram lora load

---
 utils.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/utils.py b/utils.py
index 2b6234b..5632151 100644
--- a/utils.py
+++ b/utils.py
@@ -58,6 +58,8 @@ def apply_lora(model, device_to, transformer_load_device, params_to_keep=None, d
                 name = name.replace("._orig_mod.", ".") # torch compiled modules have this prefix
                 if low_mem_load:
                     dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype
+                    if "modulation" in name:
+                        dtype_to_use = torch.float32
                     if name.startswith("diffusion_model."):
                         name_no_prefix = name[len("diffusion_model."):]
                     key = "{}.{}".format(name_no_prefix, param)
@@ -72,7 +74,9 @@ def apply_lora(model, device_to, transformer_load_device, params_to_keep=None, d
         if low_mem_load:
             for name, param in model.model.diffusion_model.named_parameters():
                 if param.device != transformer_load_device:
-                    #print("param.device", param.device)
+                    dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype
+                    if "modulation" in name:
+                        dtype_to_use = torch.float32
                     set_module_tensor_to_device(model.model.diffusion_model, name, device=transformer_load_device, dtype=dtype_to_use, value=state_dict[name])
         return model