From 4bf78e53166fc1562931fe6dca73b85f54f68bea Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 21 Jul 2025 15:25:07 +0300
Subject: [PATCH] Possibly reduce LoRA loading memory usage

---
 utils.py | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/utils.py b/utils.py
index 6f4e273..ff178c2 100644
--- a/utils.py
+++ b/utils.py
@@ -2,7 +2,12 @@ import importlib.metadata
 import torch
 import logging
 from tqdm import tqdm
-from comfy.utils import ProgressBar
+import types
+from comfy.utils import ProgressBar, copy_to_param, set_attr_param
+from comfy.model_patcher import get_key_weight, string_to_seed
+from comfy.lora import calculate_weight
+from comfy.model_management import cast_to_device
+from comfy.float import stochastic_rounding
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 log = logging.getLogger(__name__)
 
@@ -37,7 +42,32 @@ def get_tensor_memory(tensor):
     memory_bytes = tensor.element_size() * tensor.nelement()
     return f"{memory_bytes / (1024 * 1024):.2f} MB"
 
+def patch_weight_to_device(self, key, device_to=None, inplace_update=False):
+    if key not in self.patches:
+        return
+
+    weight, set_func, convert_func = get_key_weight(self.model, key)
+    inplace_update = self.weight_inplace_update or inplace_update
+
+    if device_to is not None:
+        temp_weight = cast_to_device(weight, device_to, torch.float32, copy=True)
+    else:
+        temp_weight = weight.to(torch.float32, copy=True)
+    if convert_func is not None:
+        temp_weight = convert_func(temp_weight, inplace=True)
+
+    out_weight = calculate_weight(self.patches[key], temp_weight, key)
+    if set_func is None:
+        out_weight = stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))
+        if inplace_update:
+            copy_to_param(self.model, key, out_weight)
+        else:
+            set_attr_param(self.model, key, out_weight)
+    else:
+        set_func(out_weight, inplace_update=inplace_update, seed=string_to_seed(key))
+
 def apply_lora(model, device_to, transformer_load_device, params_to_keep=None, dtype=None, base_dtype=None, state_dict=None, low_mem_load=False):
+        model.patch_weight_to_device = types.MethodType(patch_weight_to_device, model)
         to_load = []
         for n, m in model.model.named_modules():
             params = []
@@ -77,6 +107,8 @@ def apply_lora(model, device_to, transformer_load_device, params_to_keep=None, d
                     model.patch_weight_to_device("{}.{}".format(name, param), device_to=device_to, inplace_update=True)
                 else:
                     model.patch_weight_to_device("{}.{}".format(name, param), device_to=device_to)
+                    model.backup["{}.{}".format(name, param)] = None
+                    set_module_tensor_to_device(m, param, device=transformer_load_device)
                 if low_mem_load:
                     try:
                         set_module_tensor_to_device(model.model.diffusion_model, key, device=transformer_load_device, dtype=dtype_to_use, value=model.model.diffusion_model.state_dict()[key])