From c1693d0f5551f7f0b79ee9ad930c5b14869eb934 Mon Sep 17 00:00:00 2001 From: blepping Date: Sun, 25 Aug 2024 13:18:35 -0600 Subject: [PATCH] Use a GGUF array for orig shape metadata --- nodes.py | 9 +++------ tools/convert.py | 5 +---- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/nodes.py b/nodes.py index a31f778..0e61662 100644 --- a/nodes.py +++ b/nodes.py @@ -28,12 +28,9 @@ def gguf_sd_loader_get_orig_shape(reader, tensor_name): 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(","))) + if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32: + raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}") + return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data)) def gguf_sd_loader(path): """ diff --git a/tools/convert.py b/tools/convert.py index 99f4989..a49ef1c 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -156,10 +156,7 @@ def handle_tensors(args, writer, state_dict): ): orig_shape = data.shape data = data.reshape(n_params // 256, 256) - writer.add_string( - f"comfy.gguf.orig_shape.{key}", - ", ".join(str(dim) for dim in orig_shape), - ) + writer.add_array(f"comfy.gguf.orig_shape.{key}", tuple(int(dim) for dim in orig_shape)) try: data = gguf.quants.quantize(data, data_qtype)