Fix low vram lora load

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