Split orig shape metadata into individual fields

This commit is contained in:
blepping
2024-08-25 12:00:28 -06:00
parent 7413de86e6
commit 0aec56ec4c
3 changed files with 32 additions and 17 deletions
+14 -8
View File
@@ -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}:
+5 -9
View File
@@ -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
+13
View File
@@ -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