From 2e5e599b540a20a469525d85f83ea1fd140edec1 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Tue, 20 Aug 2024 02:23:37 +0200
Subject: [PATCH] CLIP Loaders

---
 nodes.py | 122 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 ops.py   |   2 +-
 2 files changed, 122 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index cf97ff9..64c010a 100644
--- a/nodes.py
+++ b/nodes.py
@@ -12,12 +12,17 @@ import comfy.model_patcher
 import folder_paths
 
 from .ops import GGMLTensor, GGMLOps
+from .dequant import dequantize_tensor
 
-# Add a custom key for files ending in .gguf
+# Add a custom keys for files ending in .gguf
 if "unet_gguf" not in folder_paths.folder_names_and_paths:
     orig = folder_paths.folder_names_and_paths.get("diffusion_models", folder_paths.folder_names_and_paths.get("unet", [[], set()]))
     folder_paths.folder_names_and_paths["unet_gguf"] = (orig[0], {".gguf"})
 
+if "clip_gguf" not in folder_paths.folder_names_and_paths:
+    orig = folder_paths.folder_names_and_paths.get("clip", [[], set()])
+    folder_paths.folder_names_and_paths["clip_gguf"] = (orig[0], {".gguf"})
+
 def gguf_sd_loader(path):
     """
     Read state dict as fake tensors
@@ -42,6 +47,41 @@ def gguf_sd_loader(path):
     print("\n")
     return sd
 
+# for remapping llama.cpp -> original key names
+clip_sd_map = {
+    "enc.": "encoder.",
+    ".blk.": ".block.",
+    "token_embd": "shared",
+    "output_norm": "final_layer_norm",
+    "attn_q": "layer.0.SelfAttention.q",
+    "attn_k": "layer.0.SelfAttention.k",
+    "attn_v": "layer.0.SelfAttention.v",
+    "attn_o": "layer.0.SelfAttention.o",
+    "attn_norm": "layer.0.layer_norm",
+    "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias",
+    "ffn_up": "layer.1.DenseReluDense.wi_1",
+    "ffn_down": "layer.1.DenseReluDense.wo",
+    "ffn_gate": "layer.1.DenseReluDense.wi_0",
+    "ffn_norm": "layer.1.layer_norm",
+}
+# weights that should be dequantized on load
+clip_sd_dequant = {
+    "shared.weight",
+}
+
+def gguf_clip_loader(path):
+    raw_sd = gguf_sd_loader(path)
+    assert "enc.blk.23.ffn_up.weight" in raw_sd, "Invalid Text Encoder!"
+    sd = {}
+    for k,v in raw_sd.items():
+        for s,d in clip_sd_map.items():
+            k = k.replace(s,d)
+        if k in clip_sd_dequant:
+            v = dequantize_tensor(v, torch.float32).to(torch.float16)
+            v = GGMLTensor(v, tensor_type=gguf.GGMLQuantizationType.F16, tensor_shape=v.shape)
+        sd[k] = v
+    return sd
+
 # TODO: Temporary fix for now
 class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
     def calculate_weight(self, patches, weight, key):
@@ -96,6 +136,86 @@ class UnetLoaderGGUF:
         model = GGUFModelPatcher.clone(model)
         return (model,)
 
+clip_name_dict = {
+    "stable_diffusion": comfy.sd.CLIPType.STABLE_DIFFUSION,
+    "stable_cascade": comfy.sd.CLIPType.STABLE_CASCADE,
+    "stable_audio": comfy.sd.CLIPType.STABLE_AUDIO,
+    "sdxl": comfy.sd.CLIPType.STABLE_DIFFUSION,
+    "sd3": comfy.sd.CLIPType.SD3,
+    "flux": comfy.sd.CLIPType.FLUX,
+}
+
+class CLIPLoaderGGUF:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip_name": (s.get_filename_list(),),
+                "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio"],),
+            }
+        }
+
+    RETURN_TYPES = ("CLIP",)
+    FUNCTION = "load_clip"
+    CATEGORY = "bootleg"
+    TITLE = "CLIPLoader (GGUF)"
+
+    @classmethod
+    def get_filename_list(s):
+        files = []
+        files += folder_paths.get_filename_list("clip")
+        files += folder_paths.get_filename_list("clip_gguf")
+        return sorted(files)
+
+    def load_data(self, ckpt_paths):
+        clip_data = []
+        for p in ckpt_paths:
+            if p.endswith(".gguf"):
+                clip_data.append(gguf_clip_loader(p))
+            else:
+                sd = comfy.utils.load_torch_file(p, safe_load=True)
+                clip_data.append(
+                    {k:GGMLTensor(v, tensor_type=gguf.GGMLQuantizationType.F16, tensor_shape=v.shape) for k,v in sd.items()}
+                )
+        return clip_data
+
+    def load_patcher(self, clip_paths, clip_type, clip_data):
+        clip = comfy.sd.load_text_encoder_state_dicts(
+            clip_type = clip_type,
+            state_dicts = clip_data,
+            model_options = {"custom_operations": GGMLOps},
+            embedding_directory = folder_paths.get_folder_paths("embeddings"),
+        )
+        clip.patcher = GGUFModelPatcher.clone(clip.patcher)
+        return clip
+
+    def load_clip(self, clip_name, type="stable_diffusion"):
+        clip_path = folder_paths.get_full_path("clip", clip_name)
+        clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION)
+        return (self.load_patcher([clip_path], clip_type, self.load_data([clip_path])),)
+
+class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip_name1": (s.get_filename_list(), ),
+                "clip_name2": (s.get_filename_list(), ),
+                "type": (["sdxl", "sd3", "flux"], ),
+            }
+        }
+
+    TITLE = "DualCLIPLoader (GGUF)"
+
+    def load_clip(self, clip_name1, clip_name2, type):
+        clip_path1 = folder_paths.get_full_path("clip", clip_name1)
+        clip_path2 = folder_paths.get_full_path("clip", clip_name2)
+        clip_paths = [clip_path1, clip_path2]
+        clip_type = clip_name_dict.get(type, comfy.sd.CLIPType.STABLE_DIFFUSION)
+        return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
+
 NODE_CLASS_MAPPINGS = {
     "UnetLoaderGGUF": UnetLoaderGGUF,
+    "CLIPLoaderGGUF": CLIPLoaderGGUF,
+    "DualCLIPLoaderGGUF": DualCLIPLoaderGGUF,
 }
diff --git a/ops.py b/ops.py
index 4358cce..5030e0d 100644
--- a/ops.py
+++ b/ops.py
@@ -119,7 +119,7 @@ class GGMLLayer(torch.nn.Module):
         bias = self.get_weight(self.bias, dtype)
         return (weight, bias)
 
-class GGMLOps(comfy.ops.disable_weight_init):
+class GGMLOps(comfy.ops.manual_cast):
     """
     Dequantize weights on the fly before doing the compute
     """