From 1c225ddb408327e267c241ede7450fb5d2b1bd50 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Tue, 15 Apr 2025 01:33:10 +0200 Subject: [PATCH] Initial auto convert script --- tools/tool_auto.py | 363 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 363 insertions(+) create mode 100644 tools/tool_auto.py diff --git a/tools/tool_auto.py b/tools/tool_auto.py new file mode 100644 index 0000000..7e40331 --- /dev/null +++ b/tools/tool_auto.py @@ -0,0 +1,363 @@ +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) +import os +import re +import sys +import time +import torch +import logging +import argparse +import subprocess +import huggingface_hub as hf + +logging.getLogger().setLevel(logging.DEBUG) + +qtypes =[ + # "F16", "BF16", + "Q8_0", "Q6_K", + "Q5_K_M", "Q5_K_S", "Q5_1", "Q5_0", + "Q4_K_M", "Q4_K_S", "Q4_1", "Q4_0", + "Q3_K_M", "Q3_K_S", "Q2_K" +] + +dtype_dict = { + "F32": torch.float32, + "F16": torch.float16, + "BF16": torch.bfloat16, + "F8_E4M3": getattr(torch, "float8_e4m3fn", "_invalid"), + "F8_E5M2": getattr(torch, "float8_e5m2", "_invalid"), +} + +# this is pretty jank but I want to be able to run it on a blank instance w/o setup +terraform_dict = { + "repo": "city96/ComfyUI-GGUF", + "target": "auto_convert", + "lcpp_repo": "ggerganov/llama.cpp", + "lcpp_target": "tags/b3962", +} + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--src", required=True, help="Source model file or huggingface repo name") + parser.add_argument("--quants", nargs="+", choices=["all", "base", *qtypes], default=["Q8_0"]) + parser.add_argument("--output-dir", default=None, help="Location for output files, defaults to current dir or ComfyUI model dir.") + parser.add_argument("--temp-dir", default=None, help="Location for temp files, defaults to [output_dir]/tmp") + parser.add_argument("--force-update", action="store_true", help="Force update & rebuild entire quantization stack.") + + args = parser.parse_args() + if args.output_dir is None: + args.output_dir = get_output_dir() + if args.temp_dir is None: + args.temp_dir = os.path.join(args.output_dir, "tmp") + + if os.path.isdir(args.temp_dir) and len(os.listdir(args.temp_dir)) > 0: + raise OSError("Output temp folder not empty!") + + if "all" in args.quants: + args.quants = ["base", *qtypes] + + return args + +def run_cmd(*args, log_error=False): + logging.debug(f"cmd: {args}") + try: + log = subprocess.run(args, capture_output=True, text=True) + except Exception as e: + logging.warning(f"{args[0]}, {e}") + return -1 + if log.returncode != 0 and log_error: + logging.warning(f"{args[0]}: {log.stdout} {log.stderr}") + else: + logging.debug(f"{args[0]}: {repr(log.stdout)} {repr(log.stderr.strip())} RET:{log.returncode}") + return log.returncode + +def setup_utils(force_update=False): + # get ComfyUI-GGUF if missing, then compile patched llama.cpp if required + root = os.path.dirname(os.path.abspath(__file__)) + root = os.path.normpath(root) + + if os.path.split(root)[1] != "tools": + cg_dir = os.path.join(root, "ComfyUI-GGUF") + if not os.path.isdir(cg_dir): + logging.warning(f"Running outside tools folder! Cloning to {cg_dir}") + run_cmd("git", "clone", f"https://github.com/{terraform_dict['repo']}", cg_dir) + need_update = True + else: + need_update = False + + if force_update or need_update: + if terraform_dict['target']: + logging.info(f"Attemtping to check out ComfyUI-GGUF branch {terraform_dict['target']}") + run_cmd("git", "-C", cg_dir, "checkout", terraform_dict['target']) + + logging.info("Attemtping to git pull ComfyUI-GGUF to latest") + run_cmd("git", "-C", cg_dir, "pull") + + tools_dir = os.path.join(root, "ComfyUI-GGUF", "tools") + sys.path.append(tools_dir) # to make import(s) work + else: + # TODO: Git pull here too? + logging.warning(f"Assuming latest ComfyUI-GGUF. Please git pull & check out branch {terraform_dict['target']} manually!") + tools_dir = root + + if not os.path.isdir(tools_dir): + raise OSError(f"Can't find tools subfoder in ComfyUI-GGUF at {tools_dir}") + + convert_path = os.path.join(tools_dir, "convert.py") + if not os.path.isfile(convert_path): + raise OSError(f"Cannot find convert.py at location: {convert_path}") + + lcpp_path = os.path.join(root, "llama.cpp.auto") # avoid messing with regular dir + if not os.path.isdir(lcpp_path): + logging.info(f"Attemtping to clone llama.cpp repo to {lcpp_path}") + run_cmd("git", "clone", f"https://github.com/{terraform_dict['lcpp_repo']}", lcpp_path) + need_update = True + else: + need_update = False + + if force_update or need_update: + # TODO: check reflog and/or git reset before checkout? + logging.info(f"Attemtping to check out llama.cpp target {terraform_dict['lcpp_target']}") + run_cmd("git", "-C", lcpp_path, "checkout", terraform_dict['lcpp_target']) + + # TODO: git reset before patch? + patch_path = os.path.join(tools_dir, "lcpp.patch") + # patch (probably) has wrong file endings: + logging.info("Converting patch file endings") + with open(patch_path, "rb") as file: + content = file.read().replace(b"\r\n", b"\n") + with open(patch_path, "wb") as file: + file.write(content) + + if run_cmd("git", "-C", lcpp_path, "apply", "--check", "-R", patch_path) != 0: + logging.info("Attemtping to apply patch to llama.cpp repo") + run_cmd("git", "-C", lcpp_path, "apply", patch_path) + else: + logging.info("Patch already applied") + + # using cmake here as llama.cpp switched to it completely for new versions + bin_name = "llama-quantize.exe" if os.name == 'nt' else "llama.quantize" + bin_path = os.path.join(lcpp_path, "build", "bin", "debug", bin_name) + if not os.path.isfile(bin_path) or force_update or need_update: + if run_cmd("cmake", "--version") != 0: + raise RuntimeError("Can't find cmake! Make sure you have a working build environment set up") + + build_path = os.path.join(lcpp_path, "build") + os.makedirs(build_path, exist_ok=True) + logging.info("Attempting to build llama.cpp binary from source") + run_cmd("cmake", "-B", build_path, lcpp_path) + run_cmd("cmake", "--build", build_path, "--config", "Debug", "-j4", "--target", "llama-quantize") + if not os.path.isfile(bin_path): + raise RuntimeError("Build failed! Rerun with --debug to see error log.") + else: + logging.info("Binary already present") + + return bin_path + +def get_output_dir(): + root = os.path.dirname(os.path.abspath(__file__)) + root = os.path.normpath(root) + split = os.path.split(root) + while split[1]: + if split[1] == "ComfyUI": + if os.path.isdir(os.path.join(*split, "models", "unet")): # new + root = os.path.join(*split, "models", "unet", "gguf") + logging.info(f"Found ComfyUI, using model folder: {root}") + return root + + if os.path.isdir(os.path.join(*split, "models", "diffusion_models")): # old + root = os.path.join(*split, "models", "diffusion_models", "gguf") + logging.info(f"Found ComfyUI, using model folder: {root}") + return root + + logging.info("Found ComfyUI, but can't find model folder") + break + + split = os.path.split(split[0]) + + root = os.path.join(root, "models") + logging.info(f"Defaulting to [script dir]/models: {root}") + return root + +def get_hf_fake_sd(repo, path, device=torch.device("meta")): + sd = {} + meta = hf.parse_safetensors_file_metadata(repo, path) + for key, raw in meta.tensors.items(): + shape = tuple(raw.shape) + dtype = dtype_dict.get(raw.dtype, torch.float32) + sd[key] = torch.zeros(shape, dtype=dtype, device=device) + return sd + +def get_hf_file_arch(repo, path): + pattern = r'(\d+)-of-(\d+)' + match = re.search(pattern, path) + + if match: + # we need to load it as multipart + if int(match.group(1)) != 1: + return None + sd = {} + for k in range(int(match.group(2))): + shard_path = path.replace(match.group(1), f"{k+1:0{len(match.group(1))}}") + sd.update(get_hf_fake_sd(repo, shard_path)) + else: + sd = get_hf_fake_sd(repo, path) + + # this should raise an error on failure + sd = strip_prefix(sd) + model_arch = detect_arch(sd) + + # this is for SDXL and SD1.5, I want to overhaul this logic to match sd.cpp eventually + assert not model_arch.shape_fix, "Model uses shape fix (SDXL/SD1) - unsupported for now." + return model_arch.arch + +def get_hf_valid_files(repo): + # TODO: probably tweak this? + MIN_SIZE_GB = 1 + VALID_SRC_EXTS = [".safetensors", ] # ".pt", ".ckpt", ] + meta = hf.model_info(repo, files_metadata=True) + + valid = {} + for file in meta.siblings: + path = file.rfilename + fname = os.path.basename(path) + name, ext = os.path.splitext(fname) + + if ext.lower() not in VALID_SRC_EXTS: + logging.debug(f"Invalid ext: {path} {ext}") + continue + + if file.size / (1024 ** 3) < MIN_SIZE_GB: + logging.debug(f"File too small: {path} {file.size}") + continue + + try: + arch = get_hf_file_arch(repo, path) + except Exception as e: + logging.warning(f"Arch detect fail: {e} ({path})") + else: + if arch is not None: + valid[path] = arch + logging.info(f"Found '{arch}' model at path {path}") + return valid + +def make_base_quant(src, output_dir, temp_dir, final=True): + name, ext = os.path.splitext(os.path.basename(src)) + if ext == ".gguf": + logging.info("Input file already in gguf, assuming base quant") + return None, src, None + + name = name.lower() # uncomment to preserve case in all quants + dst_tmp = os.path.join(temp_dir, f"{name}-{{ftype}}.gguf") # ftype is filled in by convert.py + + tmp_path, model_arch, fix_path = convert_file(src, dst_tmp, interact=False, overwrite=False) + dst_path = os.path.join(output_dir, os.path.basename(tmp_path)) + if os.path.isfile(dst_path): + raise OSError(f"Output already exists! Clear folder? {dst_path}") + + if fix_path is not None and os.path.isfile(fix_path): + quant_source = tmp_path + if final: + apply_5d_fix(tmp_path, dst_path, fix=fix_path, overwrite=False) + else: + dst_path = None + else: + fix_path = None + if final: + os.rename(tmp_path, dst_path) + quant_source = dst_path + else: + dst_path = None + quant_source = tmp_path + + return dst_path, quant_source, fix_path + +def make_quant(src, output_dir, temp_dir, qtype, quantize_binary, fix_path=None): + name, ext = os.path.splitext(os.path.basename(src)) + assert ext.lower() == ".gguf", "Invalid input file" + + src_qtext = [x for x in ["-F32.gguf", "-F16.gguf", "-BF16.gguf"] if x in src] + if len(src_qtext) == 1: + tmp_path = os.path.join( + temp_dir, + os.path.basename(src).replace(src_qtext[0], f"-{qtype.upper()}.gguf") + ) + else: + tmp_path = os.path.join( + temp_dir, + f"{name}-{qtype.upper()}.gguf" + ) + tmp_path = os.path.abspath(tmp_path) + dst_path = os.path.join(output_dir, os.path.basename(tmp_path)) + if os.path.isfile(dst_path): + raise OSError("Output already exists! Clear folder?") + + r = run_cmd(quantize_binary, src, tmp_path, qtype, log_error=True) + time.sleep(2) # leave time for file sync? + if r != 0: + raise RuntimeError(f"Quantization failed with error code {r}") + + if fix_path is not None: + apply_5d_fix(tmp_path, dst_path, fix=fix_path, overwrite=False) + if os.path.isfile(dst_path) and os.path.isfile(tmp_path): + os.remove(tmp_path) + else: + os.rename(tmp_path, dst_path) + + return dst_path + +if __name__ == "__main__": + args = get_args() + os.makedirs(args.output_dir, exist_ok=True) + os.makedirs(args.temp_dir, exist_ok=True) + quantize_binary = setup_utils(args.force_update) + + try: + from convert import detect_arch, strip_prefix, convert_file + from fix_5d_tensors import apply_5d_fix + except [ImportError, ModuleNotFoundError] as e: + raise ImportError(f"Can't import required utils: {e}") + + if not os.path.isfile(args.src): + # huggingface repo. TODO: file choice + if len(args.src.split("/")) != "1": + raise OSError(f"Invalid huggingface repo or model path {args.src}") + raise NotImplementedError("HF not yet supported") + # download then set to temp file + # hf_repo = "Lightricks/LTX-Video" # "fal/AuraFlow-v0.3" + # get_hf_valid_files(hf_repo) + # args.src = ... + + out_files = [] + + base_quant, quant_source, fix_path = make_base_quant( + args.src, + args.output_dir, + args.temp_dir, + final=("base" in args.quants) + ) + if "base" in args.quants: + args.quants = [x for x in args.quants if x not in ["base"]] + if base_quant is not None: + out_files.append(base_quant) + + for qtype in args.quants: + out_files.append(make_quant( + quant_source, + args.output_dir, + args.temp_dir, + qtype, + quantize_binary, + fix_path, + )) + + if fix_path is not None and os.path.isfile(fix_path): + os.remove(fix_path) + + if base_quant != quant_source: + # make sure our quant source is in the temp folder before removing it + cc = os.path.commonpath([os.path.normpath(quant_source), os.path.normpath(args.temp_dir)]) + if cc == os.path.normpath(args.temp_dir): + os.remove(quant_source) + + out_file_str = '\n'.join(out_files) + logging.info(f"Output file(s): {out_file_str}")