From 22e2c885f00a5ba8a0acd439a322e16c6397911a Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sat, 17 Aug 2024 07:38:55 +0200
Subject: [PATCH] Async patch load with GPU support

---
 nodes.py | 23 ++++++++++++++++++++++-
 ops.py   | 25 +++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index 7a720b5..12b9212 100644
--- a/nodes.py
+++ b/nodes.py
@@ -76,6 +76,7 @@ class LoraLoaderGGUFModelOnly:
                  "model": ("MODEL",),
                  "lora_name": (folder_paths.get_filename_list("loras"), ),
                  "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
+                 "lora_on_gpu": ("BOOLEAN", {"default": True}),
             }
         }
 
@@ -84,7 +85,7 @@ class LoraLoaderGGUFModelOnly:
     CATEGORY = "bootleg"
     TITLE = "LoRA Loader Model Only (GGUF)"
 
-    def load_lora(self, model, lora_name, strength_model):
+    def load_lora(self, model, lora_name, strength_model, lora_on_gpu):
         if strength_model == 0:
             return (model,)
 
@@ -102,6 +103,26 @@ class LoraLoaderGGUFModelOnly:
             lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
             self.loaded_lora = (lora_path, lora)
 
+        # This is a super hacky method to get it to load alongside the model
+        param_dict_name = f"lora_state_dict"#_{lora_name}".replace('.', '_')
+        target_model = model.model.diffusion_model
+        if lora_on_gpu:
+            print(f"WARNING! Loading the LoRA to GPU memory may cause OOM issues or memory leaks!")
+            temp = {}
+            params = torch.nn.ParameterDict()
+            for k,v in lora.items():
+                param = torch.nn.Parameter(v, requires_grad=False)
+                temp[k] = param
+                params[k.replace('.', '_')] = param
+            params.to(model.load_device, dtype=torch.float16)
+            target_model.register_module(param_dict_name, params)
+            lora = temp
+        else:
+            params = getattr(target_model, param_dict_name, None)
+            if params:
+                params.to("cpu", dtype=torch.float16)
+                delattr(target_model, param_dict_name)
+
         model_lora = load_lora_gguf(model, lora, strength_model)
         return (model_lora,)
 
diff --git a/ops.py b/ops.py
index 6e331de..f52faf1 100644
--- a/ops.py
+++ b/ops.py
@@ -83,12 +83,33 @@ class GGMLLayer(torch.nn.Module):
         super()._apply(fn)
         return self
 
+    def move_patch_to_cuda(self, item, device):
+        if isinstance(item, torch.Tensor):
+            return item.to(device, non_blocking=True)
+        elif isinstance(item, tuple):
+            return tuple(self.move_patch_to_cuda(x, device) for x in item)
+        elif isinstance(item, list):
+            return [self.move_patch_to_cuda(x, device) for x in item]
+        else:
+            return item
+
     def get_weight(self, tensor, dtype):
         if tensor is None:
             return
+
+        # consolidate and load patches to GPU in async
+        patch_list = []
+        device = tensor.device
+        t_move = lambda x: x.to(device) if torch.is_tensor(x) else x
+        for function, patches, _ in getattr(tensor, "patches", []):
+            patch_list += self.move_patch_to_cuda(patches, device)
+
+        # dequantize tensor while patches load
         weight = dequantize_tensor(tensor, dtype)
-        for function, patches, key in getattr(tensor, "patches", []):
-            weight = function(patches, weight, key)
+
+        # apply patches
+        if patch_list:
+            weight = function(patch_list, weight, "dequant.name.unknown")
         return weight
 
     def get_weights(self, dtype=torch.float16):