Don't dequantize token embed on load

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
 

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)