From d99b2d931fabe97c963e091c5ac24663d9961900 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 14 Mar 2025 16:56:48 +0100
Subject: [PATCH] force dequant on umT5 to prevent OOM

not the best solution but that tensor alone needs ~3GB of runtime memory to dequant
---
 loader.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/loader.py b/loader.py
index ea3a0e1..18c1af1 100644
--- a/loader.py
+++ b/loader.py
@@ -3,7 +3,7 @@ import torch
 import gguf
 
 from .ops import GGMLTensor
-from .dequant import is_quantized
+from .dequant import is_quantized, dequantize_tensor
 
 IMG_ARCH_LIST = {"flux", "sd1", "sdxl", "sd3", "aura", "ltxv", "hyvid", "wan"}
 TXT_ARCH_LIST = {"t5", "t5encoder", "llama"}
@@ -230,6 +230,9 @@ def gguf_clip_loader(path):
         if temb_key in sd and sd[temb_key].shape == (256384, 4096):
             # non-standard Comfy-Org tokenizer
             sd["spiece_model"] = gguf_tokenizer_loader(path, sd[temb_key].shape)
+            # TODO: dequantizing token embed here is janky but otherwise we OOM due to tensor being massive.
+            print(f"Dequantizing {temb_key} to prevent runtime OOM.")
+            sd[temb_key] = dequantize_tensor(sd[temb_key], dtype=torch.float16)
         sd = sd_map_replace(sd, T5_SD_MAP)
     elif arch in {"llama"}:
         temb_key = "token_embd.weight"