From 6618801f1c1ee6c4203e4a76b87d17f3a2a3d3eb Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 30 Aug 2024 06:25:19 +0200
Subject: [PATCH] Don't dequantize token embed on load

---
 nodes.py |  7 -------
 ops.py   | 12 +++++++++++-
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/nodes.py b/nodes.py
index 70823f4..9b7bede 100644
--- a/nodes.py
+++ b/nodes.py
@@ -78,10 +78,6 @@ clip_sd_map = {
     "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)
@@ -90,9 +86,6 @@ def gguf_clip_loader(path):
     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
 
diff --git a/ops.py b/ops.py
index 43f6299..bc17c3a 100644
--- a/ops.py
+++ b/ops.py
@@ -109,7 +109,7 @@ class GGMLLayer(torch.nn.Module):
     def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
         if input is not None:
             if dtype is None:
-                dtype = input.dtype
+                dtype = getattr(input, "dtype", torch.float32)
             if bias_dtype is None:
                 bias_dtype = dtype
             if device is None:
@@ -139,6 +139,16 @@ class GGMLOps(comfy.ops.manual_cast):
             weight, bias = self.cast_bias_weight(input)
             return self._conv_forward(input, weight, bias)
 
+    class Embedding(GGMLLayer, comfy.ops.manual_cast.Embedding):
+        def forward_comfy_cast_weights(self, input, out_dtype=None):
+            output_dtype = out_dtype
+            if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
+                out_dtype = None
+            weight, bias = self.cast_bias_weight(self, device=input.device, dtype=out_dtype)
+            return torch.nn.functional.embedding(
+                input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse
+            ).to(dtype=output_dtype)
+
 def move_patch_to_cuda(item, device):
     if isinstance(item, torch.Tensor):
         return item.to(device, non_blocking=True)