nodes/city96-ComfyUI-GGUF
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

K quant patch / instructions

Browse Source 
#11
This commit is contained in:
City
2024-08-21 06:31:40 +02:00
parent 1ebce77a44
commit 8d1fbbc456
4 changed files with 255 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/convert.py
+45 -57
View File
@@ -1,5 +1,6 @@
# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
import os
import sys
import torch
import numpy as np
import gguf # This needs to be the llama.cpp one specifically!
@@ -9,27 +10,13 @@ from tqdm import tqdm
from safetensors.torch import load_file
def parse_args():
parser = argparse.ArgumentParser(description="Generate GGUF files from single SD ckpt")
parser = argparse.ArgumentParser(description="Generate F16 GGUF files from single UNET")
parser.add_argument("--src", required=True, help="Source model ckpt file.")
parser.add_argument("--dst", help="Output unet gguf file.")
parser.add_argument("--qtype", default="F16", help="Quant type [default: f16]")
args = parser.parse_args()
if not os.path.isfile(args.src):
parser.error("No input provided!")
if args.dst is None:
args.dst = os.path.splitext(args.src)[0] + f"_{args.qtype}.gguf"
args.dst = os.path.basename(args.dst)
if os.path.isfile(args.dst):
input("Output exists enter to continue or ctrl+c to abort!")
try:
args.ftype = getattr(gguf.LlamaFileType, f"MOSTLY_{args.qtype}")
args.qtype = getattr(gguf.GGMLQuantizationType, args.qtype)
except AttributeError:
parser.error(f"Unknown quant/file type {args.qtype}")
return args
@@ -39,10 +26,21 @@ def load_state_dict(path):
state_dict = state_dict.get("model", state_dict)
else:
state_dict = load_file(path)
return state_dict
# only keep unet with no prefix!
sd = {}
has_prefix = any(["model.diffusion_model." in x for x in state_dict.keys()])
for k, v in state_dict.items():
if has_prefix and "model.diffusion_model." not in k:
continue
if has_prefix:
k = k.replace("model.diffusion_model.", "")
sd[k] = v
def load_model(args):
state_dict = load_state_dict(args.src)
return sd
def load_model(path):
state_dict = load_state_dict(path)
# from ComfyUI model detection
if "transformer_blocks.0.attn.norm_added_k.weight" in state_dict:
@@ -64,34 +62,31 @@ def load_model(args):
writer = gguf.GGUFWriter(path=None, arch=arch)
return (writer, state_dict)
def handle_metadata(args, writer, state_dict):
# TODO: actual metadata
writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
writer.add_file_type(args.ftype)
def handle_tensors(args, writer, state_dict):
# TODO list:
# - do something about this being awful and hacky
max_name_len = max([len(s) for s in state_dict.keys()]) + 4
for key, data in tqdm(state_dict.items()):
if data.dtype == torch.bfloat16:
data = data.to(torch.float32)
data = data.numpy()
old_dtype = data.dtype
if data.dtype == torch.bfloat16:
data = data.to(torch.float32).numpy()
else:
data = data.numpy()
n_dims = len(data.shape)
data_qtype = args.qtype
data_shape = data.shape
data_qtype = getattr(
gguf.GGMLQuantizationType,
"BF16" if old_dtype == torch.bfloat16 else "F16"
)
# get number of parameters (AKA elements) in this tensor
n_params = 1
for dim_size in data_shape:
n_params *= dim_size
fallback = gguf.GGMLQuantizationType.F16
# keys to keep as max precision
blacklist = [
"time_embedding.",
@@ -105,7 +100,7 @@ def handle_tensors(args, writer, state_dict):
]
if any([x in key for x in blacklist]) and ".weight" in key:
data_qtype = fallback
data_qtype = gguf.GGMLQuantizationType.F32
if n_dims == 1:
# one-dimensional tensors should be kept in F32
@@ -118,20 +113,12 @@ def handle_tensors(args, writer, state_dict):
elif n_dims == 4:
if min(data.shape[:2]) == 4: # output tensor
data_qtype = fallback
data_qtype = gguf.GGMLQuantizationType.F16
elif data_shape[-1] == 3: # 3x3 kernel
data_qtype = fallback
data_qtype = gguf.GGMLQuantizationType.F16
elif data_shape[-1] == 1: # 1x1 kernel
#data = np.squeeze(data) # don't do this
data_qtype = fallback
# TODO: find keys to keep in higher precision(s) / qtypes
# if "time_emb_proj.weight" in key:
# data_qtype = gguf.GGMLQuantizationType.F16
# if ".to_v.weight" in key or ".to_out" in key:
# data_qtype = gguf.GGMLQuantizationType.F16
# if "ff.net" in key:
# data_qtype = gguf.GGMLQuantizationType.F16
data_qtype = gguf.GGMLQuantizationType.F16
try:
data = gguf.quants.quantize(data, data_qtype)
@@ -144,7 +131,6 @@ def handle_tensors(args, writer, state_dict):
data_qtype = gguf.GGMLQuantizationType.F16
data = gguf.quants.quantize(data, data_qtype)
assert len(key) < 64, f"Invalid key length! Cannot store in gguf file. {key}"
new_name = key # do we need to rename?
shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}"
@@ -152,23 +138,25 @@ def handle_tensors(args, writer, state_dict):
writer.add_tensor(new_name, data, raw_dtype=data_qtype)
warning = """
######################################################
The quantized file format needs more work.
Consider **not** uploading the resulting files for now
######################################################
"""
if __name__ == "__main__":
args = parse_args()
writer, state_dict = load_model(args)
handle_metadata(args, writer, state_dict)
handle_tensors(args, writer, state_dict)
path = args.src
writer, state_dict = load_model(path)
writer.write_header_to_file(path=(args.dst or "test.gguf"))
writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
if next(iter(state_dict.values())).dtype == torch.bfloat16:
out_path = f"{os.path.splitext(path)[0]}-BF16.gguf"
writer.add_file_type(gguf.LlamaFileType.MOSTLY_BF16)
else:
out_path = f"{os.path.splitext(path)[0]}-F16.gguf"
writer.add_file_type(gguf.LlamaFileType.MOSTLY_F16)
out_path = args.dst or out_path
if os.path.isfile(out_path):
input("Output exists enter to continue or ctrl+c to abort!")
handle_tensors(path, writer, state_dict)
writer.write_header_to_file(path=out_path)
writer.write_kv_data_to_file()
writer.write_tensors_to_file(progress=True)
writer.close()
print(warning)