From 4c7d55a3c3a9ada8ff6c32ae51435543a6e19aab Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Fri, 18 Aug 2023 12:52:12 -0400
Subject: [PATCH] Custom node for creating control loras.

---
 README.md              |   9 ++++
 __init__.py            |  17 +++++++
 control_lora_create.py | 106 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 132 insertions(+)
 create mode 100644 README.md
 create mode 100644 __init__.py
 create mode 100644 control_lora_create.py

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..3e6ee9f
--- /dev/null
+++ b/README.md
@@ -0,0 +1,9 @@
+# Custom nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
+
+git clone this repo to your ComfyUI/custom_nodes folder, it should look like: ComfyUI/custom_nodes/stability-ComfyUI-nodes
+
+These nodes will appear in the stability section.
+
+### ControlLoraSave
+
+This node can be used to create a Control Lora from a model and a controlnet. It will take the difference between the model weights and the controlnet weights and store that difference in Lora format.
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..c1dc126
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,17 @@
+import importlib
+import os
+
+node_list = [ #Add list of .py files containing nodes here
+    "control_lora_create",
+]
+
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+for module_name in node_list:
+    imported_module = importlib.import_module(".{}".format(module_name), __name__)
+
+    NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **control_lora_create.NODE_CLASS_MAPPINGS}
+    NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **control_lora_create.NODE_DISPLAY_NAME_MAPPINGS}
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/control_lora_create.py b/control_lora_create.py
new file mode 100644
index 0000000..923fb1b
--- /dev/null
+++ b/control_lora_create.py
@@ -0,0 +1,106 @@
+import torch
+import comfy.model_management
+import comfy.utils
+import folder_paths
+import os
+
+CLAMP_QUANTILE = 0.99
+
+def extract_lora(diff, rank):
+    conv2d = (len(diff.shape) == 4)
+    kernel_size = None if not conv2d else diff.size()[2:4]
+    conv2d_3x3 = conv2d and kernel_size != (1, 1)
+    out_dim, in_dim = diff.size()[0:2]
+    rank = min(rank, in_dim, out_dim)
+
+    if conv2d:
+        if conv2d_3x3:
+            diff = diff.flatten(start_dim=1)
+        else:
+            diff = diff.squeeze()
+
+
+    U, S, Vh = torch.linalg.svd(diff.float())
+    U = U[:, :rank]
+    S = S[:rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:rank, :]
+
+    dist = torch.cat([U.flatten(), Vh.flatten()])
+    hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+    low_val = -hi_val
+
+    U = U.clamp(low_val, hi_val)
+    Vh = Vh.clamp(low_val, hi_val)
+    if conv2d:
+        U = U.reshape(out_dim, rank, 1, 1)
+        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
+    return (U, Vh)
+
+class ControlLoraSave:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "control_net": ("CONTROL_NET",),
+                              "filename_prefix": ("STRING", {"default": "controlnet_loras/ComfyUI_control_lora"}),
+                              "rank": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 8}),
+                            },}
+    RETURN_TYPES = ()
+    FUNCTION = "save"
+    OUTPUT_NODE = True
+
+    CATEGORY = "stability/controlnet"
+
+    def save(self, model, control_net, filename_prefix, rank):
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+
+        output_sd = {}
+        prefix_key = "diffusion_model."
+        stored = set()
+
+        comfy.model_management.load_models_gpu([model])
+        f = model.model_state_dict()
+        c = control_net.control_model.state_dict()
+
+        for k in f:
+            if k.startswith(prefix_key):
+                ck = k[len(prefix_key):]
+                if ck not in c:
+                    ck = "control_model.{}".format(ck)
+                if ck in c:
+                    model_weight = f[k]
+                    if len(model_weight.shape) >= 2:
+                        diff = c[ck].float().to(model_weight.device) - model_weight.float()
+                        out = extract_lora(diff, rank)
+                        name = ck
+                        if name.endswith(".weight"):
+                            name = name[:-len(".weight")]
+                        out1_key = "{}.up".format(name)
+                        out2_key = "{}.down".format(name)
+                        output_sd[out1_key] = out[0].contiguous().half().cpu()
+                        output_sd[out2_key] = out[1].contiguous().half().cpu()
+                    else:
+                        output_sd[ck] = c[ck]
+                        print(ck, c[ck].shape)
+                    stored.add(ck)
+
+        for k in c:
+            if k not in stored:
+                output_sd[k] = c[k].half()
+        output_sd["lora_controlnet"] = torch.tensor([])
+
+        output_checkpoint = f"{filename}_{counter:05}_.safetensors"
+        output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
+
+        comfy.utils.save_torch_file(output_sd, output_checkpoint, metadata=None)
+        return {}
+
+NODE_CLASS_MAPPINGS = {
+    "ControlLoraSave": ControlLoraSave
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+}