From 624687939f1d361b59770f0a0b5bf75b7d5892e9 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Thu, 10 Apr 2025 07:12:10 +0200 Subject: [PATCH] Refactor 5D patch and save loop --- tools/convert.py | 98 ++++++++++++++++++++++++++--------------- tools/fix_5d_tensors.py | 7 ++- 2 files changed, 67 insertions(+), 38 deletions(-) diff --git a/tools/convert.py b/tools/convert.py index 52d1800..6c5669d 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -1,11 +1,10 @@ # (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) import os +import gguf import torch -import gguf # This needs to be the llama.cpp one specifically! import argparse from tqdm import tqdm - -from safetensors.torch import load_file +from safetensors.torch import load_file, save_file QUANTIZATION_THRESHOLD = 1024 REARRANGE_THRESHOLD = 512 @@ -55,14 +54,14 @@ class ModelHyVid(ModelTemplate): ) ] - # def handle_nd_tensor(self, key, data): - # # hacky but don't have any better ideas - # path = f"./fix_5d_tensors_{self.arch}.pt" - # if os.path.isfile(path): - # raise RuntimeError(f"5D tensor fix file already exists! {path}") - # fsd = {key: data} - # tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") - # torch.save(fsd, path) + def handle_nd_tensor(self, key, data): + # hacky but don't have any better ideas + path = f"./fix_5d_tensors_{self.arch}.safetensors" # TODO: somehow get a path here?? + if os.path.isfile(path): + raise RuntimeError(f"5D tensor fix file already exists! {path}") + fsd = {key: torch.from_numpy(data)} + tqdm.write(f"5D key found in state dict! Manual fix required! - {key} {data.shape}") + save_file(fsd, path) class ModelWan(ModelHyVid): arch = "wan" @@ -180,13 +179,6 @@ def load_state_dict(path): return strip_prefix(state_dict) -def load_model(path): - state_dict = load_state_dict(path) - model_arch = detect_arch(state_dict) - print(f"* Architecture detected from input: {model_arch.arch}") - writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) - return (writer, state_dict, model_arch) - def handle_tensors(writer, state_dict, model_arch): name_lengths = tuple(sorted( ((key, len(key)) for key in state_dict.keys()), @@ -212,10 +204,12 @@ def handle_tensors(writer, state_dict, model_arch): n_dims = len(data.shape) data_shape = data.shape - data_qtype = getattr( - gguf.GGMLQuantizationType, - "BF16" if old_dtype == torch.bfloat16 else "F16" - ) + if old_dtype == torch.bfloat16: + data_qtype = gguf.GGMLQuantizationType.BF16 + # elif old_dtype == torch.float32: + # data_qtype = gguf.GGMLQuantizationType.F32 + else: + data_qtype = gguf.GGMLQuantizationType.F16 # The max no. of dimensions that can be handled by the quantization code is 4 if len(data.shape) > MAX_TENSOR_DIMS: @@ -265,25 +259,57 @@ def handle_tensors(writer, state_dict, model_arch): writer.add_tensor(new_name, data, raw_dtype=data_qtype) -if __name__ == "__main__": - args = parse_args() - path = args.src - writer, state_dict, model_arch = load_model(path) +def convert_file(path, dst_path=None, interact=True, overwrite=False): + # load & run model detection logic + state_dict = load_state_dict(path) + model_arch = detect_arch(state_dict) + print(f"* Architecture detected from input: {model_arch.arch}") - 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) + # detect & set dtype for output file + dtypes = [x.dtype for x in state_dict.values()] + dtypes = {x:dtypes.count(x) for x in set(dtypes)} + main_dtype = max(dtypes, key=dtypes.get) + + if main_dtype == torch.bfloat16: + ftype_name = "BF16" + ftype_gguf = gguf.LlamaFileType.MOSTLY_BF16 + # elif main_dtype == torch.float32: + # ftype_name = "F32" + # ftype_gguf = None else: - out_path = f"{os.path.splitext(path)[0]}-F16.gguf" - writer.add_file_type(gguf.LlamaFileType.MOSTLY_F16) + ftype_name = "F16" + ftype_gguf = 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!") + if dst_path is None: + dst_path = f"{os.path.splitext(path)[0]}-{ftype_name}.gguf" + elif "{ftype}" in dst_path: # lcpp logic + dst_path = dst_path.replace("{ftype}", ftype_name) + + if os.path.isfile(dst_path) and not overwrite: + if interact: + input("Output exists enter to continue or ctrl+c to abort!") + else: + raise OSError("Output exists and overwriting is disabled!") + + # handle actual file + writer = gguf.GGUFWriter(path=None, arch=model_arch.arch) + writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + if ftype_gguf is not None: + writer.add_file_type(ftype_gguf) handle_tensors(writer, state_dict, model_arch) - writer.write_header_to_file(path=out_path) + writer.write_header_to_file(path=dst_path) writer.write_kv_data_to_file() writer.write_tensors_to_file(progress=True) writer.close() + + fix = f"./fix_5d_tensors_{model_arch.arch}.safetensors" + if os.path.isfile(fix): + print(f"\n### Warning! Fix file found at '{fix}'") + print(f" you most likely need to run 'fix_5d_tensors.py' after quantization.") + + return dst_path, model_arch + +if __name__ == "__main__": + args = parse_args() + convert_file(args.src, args.dst) diff --git a/tools/fix_5d_tensors.py b/tools/fix_5d_tensors.py index e59794e..0e61d1c 100644 --- a/tools/fix_5d_tensors.py +++ b/tools/fix_5d_tensors.py @@ -1,8 +1,10 @@ +# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0) import os import gguf import torch import argparse from tqdm import tqdm +from safetensors.torch import load_file def get_args(): parser = argparse.ArgumentParser() @@ -39,12 +41,13 @@ if __name__ == "__main__": # prep fix if args.fix is None: - args.fix = f"./fix_5d_tensors_{arch}.pt" + args.fix = f"./fix_5d_tensors_{arch}.safetensors" if not os.path.isfile(args.fix): raise OSError(f"No 5D tensor fix file: {args.fix}") - sd5d = torch.load(args.fix) + sd5d = load_file(args.fix) + sd5d = {k:v.numpy() for k,v in sd5d.items()} print("5D tensors:", sd5d.keys()) # prep output