Async patch load with GPU support

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,)