Use a GGUF array for orig shape metadata

This commit is contained in:
blepping
2024-08-25 13:18:35 -06:00
parent 0aec56ec4c
commit c1693d0f55
2 changed files with 4 additions and 10 deletions
+3 -6
View File
@@ -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):
"""
+1 -4
View File
@@ -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)