From 303af6ff756d1379802d041712b4d1eb992d6a1a Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Thu, 15 Aug 2024 08:09:59 +0200 Subject: [PATCH] Upload basic quantization code --- tools/README.md | 19 ++++++ tools/convert.py | 160 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 179 insertions(+) create mode 100644 tools/README.md create mode 100644 tools/convert.py diff --git a/tools/README.md b/tools/README.md new file mode 100644 index 0000000..ead935a --- /dev/null +++ b/tools/README.md @@ -0,0 +1,19 @@ +This needs the llama.cpp version of gguf-py to work at the moment, not the pip one as that one does not have the python quantization code yet. + +``` +git clone https://github.com/ggerganov/llama.cpp +pip install llama.cpp/gguf-py +``` + +To quantize: +``` +python convert.py --src ~/ComfyUI/models/unet/flux1-dev.safetensors --dst ~/ComfyUI/models/unet/flux1-dev-Q4_0.gguf --qtype Q4_0 +``` + +Working quant types: Q4_0, Q5_0, Q8_0, F16 + +> [!WARNING] +> Do not use the diffusers UNET for flux, it won't work, use the default checkpoint that comes with the model or convert it. + +> [!IMPORTANT] +> The model format is very much WIP. I don't recommend uploading the model files created with this method anywhere until proper metadata is added, although the key/quantization format is unlikely to change. diff --git a/tools/convert.py b/tools/convert.py new file mode 100644 index 0000000..7b44300 --- /dev/null +++ b/tools/convert.py @@ -0,0 +1,160 @@ +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import os +import torch +import numpy as np +import gguf # This needs to be the llama.cpp one specifically! +import argparse +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.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.qtype = getattr(gguf.GGMLQuantizationType, args.qtype) + except AttributeError: + parser.error(f"Unknown quant type {args.qtype}") + + return args + +def load_state_dict(path): + if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): + state_dict = torch.load(path, map_location="cpu", weights_only=True) + state_dict = state_dict.get("model", state_dict) + else: + state_dict = load_file(path) + return state_dict + +def load_model(args): + state_dict = load_state_dict(args.src) + + # from ComfyUI model detection + if "transformer_blocks.0.attn.norm_added_k.weight" in state_dict: + arch = "flux" + raise ValueError(f"The Diffusers UNET can not be used for this!") + elif "double_blocks.0.img_attn.proj.weight" in state_dict: + arch = "flux" # mmdit ...? + elif "transformer_blocks.0.attn.add_q_proj.weight" in state_dict: + arch = "sd3" + elif "down_blocks.0.downsamplers.0.conv.weight" in state_dict: + if "add_embedding.linear_1.weight" in state_dict: + arch = "sdxl" + else: + arch = "sd1" + else: + breakpoint() + raise ValueError(f"Unknown model architecture!") + + 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.qtype) # TODO: gguf.LlamaFileType + +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 + + n_dims = len(data.shape) + data_qtype = args.qtype + data_shape = data.shape + fallback = gguf.GGMLQuantizationType.F16 + + if n_dims == 1: + # these barely take up space, just use F32 / F16. + # also speeds up inference due to not dequantizing + data_qtype = fallback + elif n_dims == 4: + if min(data.shape[:2]) == 4: # output tensor + data_qtype = fallback + elif data.shape[-1] == 3: # 3x3 kernel + data_qtype = fallback + elif data.shape[-1] == 1: # 1x1 kernel + #data = np.squeeze(data) # don't do this + data_qtype = fallback + + # other keys to keep as max precision + blacklist = [ + "time_embedding.", + "add_embedding.", + "time_in.", + "txt_in.", + "vector_in.", + "img_in.", + "guidance_in.", + "final_layer.", + ] + if any([x in key for x in blacklist]) and ".weight" in key: + 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 + + try: + data = gguf.quants.quantize(data, data_qtype) + except gguf.QuantError as e: + tqdm.write(f"falling back to F16: {e}") + data_qtype = gguf.GGMLQuantizationType.F16 + data = gguf.quants.quantize(data, data_qtype) + except AttributeError as e: + tqdm.write(f"falling back to F16: {e}") + data_qtype = gguf.GGMLQuantizationType.F16 + data = gguf.quants.quantize(data, data_qtype) + + new_name = key # do we need to rename? + + shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}" + tqdm.write(f"{f'%-{max_name_len}s' % f'{new_name},'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") + + 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) + + writer.write_header_to_file(path=(args.dst or "test.gguf")) + writer.write_kv_data_to_file() + writer.write_tensors_to_file(progress=True) + writer.close() + + print(warning)