Refactor 5D patch and save loop

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)

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