SaveLora node can now save "full diff" lora format.

This isn't actually a lora format and is saving the full diff of the
weights in a format that can be used in the lora loader nodes.

comfy_extras/nodes_lora_extract.py

@@ -4,6 +4,7 @@
 import folder_paths
 import os
 import logging
+from enum import Enum
 
 CLAMP_QUANTILE = 0.99
 
@@ -38,21 +39,34 @@ def extract_lora(diff, rank):
         Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
     return (U, Vh)
 
-def calc_lora_model(model_diff, rank, prefix_model, prefix_lora, output_sd, bias_diff=False):
+class LORAType(Enum):
+    STANDARD = 0
+    FULL_DIFF = 1
+
+LORA_TYPES = {"standard": LORAType.STANDARD,
+              "full_diff": LORAType.FULL_DIFF}
+
+def calc_lora_model(model_diff, rank, prefix_model, prefix_lora, output_sd, lora_type, bias_diff=False):
     comfy.model_management.load_models_gpu([model_diff], force_patch_weights=True)
     sd = model_diff.model_state_dict(filter_prefix=prefix_model)
 
     for k in sd:
         if k.endswith(".weight"):
             weight_diff = sd[k]
-            if weight_diff.ndim < 2:
-                continue
-            try:
-                out = extract_lora(weight_diff, rank)
-                output_sd["{}{}.lora_up.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[0].contiguous().half().cpu()
-                output_sd["{}{}.lora_down.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[1].contiguous().half().cpu()
-            except:
-                logging.warning("Could not generate lora weights for key {}, is the weight difference a zero?".format(k))
+            if lora_type == LORAType.STANDARD:
+                if weight_diff.ndim < 2:
+                    if bias_diff:
+                        output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().half().cpu()
+                    continue
+                try:
+                    out = extract_lora(weight_diff, rank)
+                    output_sd["{}{}.lora_up.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[0].contiguous().half().cpu()
+                    output_sd["{}{}.lora_down.weight".format(prefix_lora, k[len(prefix_model):-7])] = out[1].contiguous().half().cpu()
+                except:
+                    logging.warning("Could not generate lora weights for key {}, is the weight difference a zero?".format(k))
+            elif lora_type == LORAType.FULL_DIFF:
+                output_sd["{}{}.diff".format(prefix_lora, k[len(prefix_model):-7])] = weight_diff.contiguous().half().cpu()
+
         elif bias_diff and k.endswith(".bias"):
             output_sd["{}{}.diff_b".format(prefix_lora, k[len(prefix_model):-5])] = sd[k].contiguous().half().cpu()
     return output_sd
@@ -65,7 +79,7 @@ def __init__(self):
     def INPUT_TYPES(s):
         return {"required": {"filename_prefix": ("STRING", {"default": "loras/ComfyUI_extracted_lora"}),
                               "rank": ("INT", {"default": 8, "min": 1, "max": 4096, "step": 1}),
-                              "lora_type": (["standard"],),
+                              "lora_type": (tuple(LORA_TYPES.keys()),),
                               "bias_diff": ("BOOLEAN", {"default": True}),
                             },
                 "optional": {"model_diff": ("MODEL",),
@@ -81,13 +95,14 @@ def save(self, filename_prefix, rank, lora_type, bias_diff, model_diff=None, tex
         if model_diff is None and text_encoder_diff is None:
             return {}
 
+        lora_type = LORA_TYPES.get(lora_type)
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
 
         output_sd = {}
         if model_diff is not None:
-            output_sd = calc_lora_model(model_diff, rank, "diffusion_model.", "diffusion_model.", output_sd, bias_diff=bias_diff)
+            output_sd = calc_lora_model(model_diff, rank, "diffusion_model.", "diffusion_model.", output_sd, lora_type, bias_diff=bias_diff)
         if text_encoder_diff is not None:
-            output_sd = calc_lora_model(text_encoder_diff.patcher, rank, "", "text_encoders.", output_sd, bias_diff=bias_diff)
+            output_sd = calc_lora_model(text_encoder_diff.patcher, rank, "", "text_encoders.", output_sd, lora_type, bias_diff=bias_diff)
 
         output_checkpoint = f"{filename}_{counter:05}_.safetensors"
         output_checkpoint = os.path.join(full_output_folder, output_checkpoint)