From 0aec56ec4c5c7a32f8fb465503d26d3c78bed685 Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Sun, 25 Aug 2024 12:00:28 -0600
Subject: [PATCH] Split orig shape metadata into individual fields

---
 nodes.py         | 22 ++++++++++++++--------
 tools/convert.py | 14 +++++---------
 tools/lcpp.patch | 13 +++++++++++++
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/nodes.py b/nodes.py
index af15bcf..a31f778 100644
--- a/nodes.py
+++ b/nodes.py
@@ -22,24 +22,30 @@ 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_get_orig_shape(reader, tensor_name):
+    field_key = f"comfy.gguf.orig_shape.{tensor_name}"
+    field = reader.get_field(field_key)
+    if field is None:
+        return None
+    # Has original shape metadata, so we try to decode it.
+    part_idx = field.data[-1]
+    part_type = field.types[0] if len(field.types) == 1 else "INVALID"
+    if part_type != gguf.GGUFValueType.STRING:
+        raise TypeError(f"Bad original shape metadata for {field_key}: Expected STRING, got {part_type}")
+    shape_str = str(field.parts[part_idx], encoding="utf-8")
+    return torch.Size(tuple(int(dim) for dim in shape_str.split(",")))
+
 def gguf_sd_loader(path):
     """
     Read state dict as fake tensors
     """
     reader = gguf.GGUFReader(path)
-    fixup_field = reader.get_field("comfy.gguf.tensor_fixup")
-    tensor_fixup = {}
-    if fixup_field is not None:
-        tensor_fixup_str=str(fixup_field.parts[fixup_field.data[-1]], encoding="utf-8")
-        for line in str(tensor_fixup_str).split("\n"):
-            key, *fixup_items = line.split(" ")
-            tensor_fixup[key] = torch.Size(tuple(int(i) for i in fixup_items))
     sd = {}
     dt = {}
     for tensor in reader.tensors:
         tensor_name = str(tensor.name)
-        shape = tensor_fixup.get(tensor_name)
         torch_tensor = torch.from_numpy(tensor.data) # mmap
+        shape = gguf_sd_loader_get_orig_shape(reader, tensor_name)
         if shape is None:
             shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
         elif tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
diff --git a/tools/convert.py b/tools/convert.py
index 44aabb9..99f4989 100644
--- a/tools/convert.py
+++ b/tools/convert.py
@@ -9,7 +9,7 @@ from safetensors.torch import load_file
 
 QUANTIZATION_THRESHOLD = 1024
 REARRANGE_THRESHOLD = 512
-MAX_TENSOR_NAME_LENGTH = 63
+MAX_TENSOR_NAME_LENGTH = 127
 
 # Tuple of arch_name, match_lists.
 # Each item in match_lists is a tuple of keys that must match.
@@ -104,7 +104,6 @@ def handle_tensors(args, writer, state_dict):
     if max_name_len > MAX_TENSOR_NAME_LENGTH:
         bad_list = ", ".join(f"{key!r} ({namelen})" for key, namelen in name_lengths if namelen > MAX_TENSOR_NAME_LENGTH)
         raise ValueError(f"Can only handle tensor names up to {MAX_TENSOR_NAME_LENGTH} characters. Tensors exceeding the limit: {bad_list}")
-    tensor_fixup = {}
     for key, data in tqdm(state_dict.items()):
         old_dtype = data.dtype
 
@@ -157,7 +156,10 @@ def handle_tensors(args, writer, state_dict):
         ):
             orig_shape = data.shape
             data = data.reshape(n_params // 256, 256)
-            tensor_fixup[key] = orig_shape
+            writer.add_string(
+                f"comfy.gguf.orig_shape.{key}",
+                ", ".join(str(dim) for dim in orig_shape),
+            )
 
         try:
             data = gguf.quants.quantize(data, data_qtype)
@@ -173,12 +175,6 @@ def handle_tensors(args, writer, state_dict):
 
         writer.add_tensor(new_name, data, raw_dtype=data_qtype)
 
-    if not tensor_fixup:
-        return
-
-    fixup_str = "\n".join(" ".join((k,)+tuple(str(dimval) for dimval in v)) for k,v in tensor_fixup.items())
-    writer.add_string("comfy.gguf.tensor_fixup", fixup_str)
-
 if __name__ == "__main__":
     args = parse_args()
     path = args.src
diff --git a/tools/lcpp.patch b/tools/lcpp.patch
index 01fe0b8..6355cb8 100644
--- a/tools/lcpp.patch
+++ b/tools/lcpp.patch
@@ -1,3 +1,16 @@
+diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
+index 1d2a3540..7d354dc6 100644
+--- a/ggml/include/ggml.h
++++ b/ggml/include/ggml.h
+@@ -230,7 +230,7 @@
+ #define GGML_MAX_CONTEXTS       64
+ #define GGML_MAX_SRC            10
+ #ifndef GGML_MAX_NAME
+-#define GGML_MAX_NAME           64
++#define GGML_MAX_NAME           128
+ #endif
+ #define GGML_MAX_OP_PARAMS      64
+ #define GGML_DEFAULT_N_THREADS  4
 diff --git a/src/llama.cpp b/src/llama.cpp
 index 5ab65ea9..6d684fef 100644
 --- a/src/llama.cpp