Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
71 lines
3.6 KiB
Python
71 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Package RE-USE directly from its original safetensors checkpoint."""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import logging
|
|
from pathlib import Path
|
|
import subprocess
|
|
|
|
import gguf
|
|
import numpy as np
|
|
from safetensors import safe_open
|
|
|
|
|
|
def main():
|
|
root = Path(__file__).resolve().parents[2]
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--type", choices=("f32", "f16", "bf16", "q8_0"), default="f32")
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
parser.add_argument("--audiocpp-gguf", type=Path, default=root / "build/debug/bin/audiocpp_gguf")
|
|
parser.add_argument("--log", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
args.log.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(args.log, mode="w"), logging.StreamHandler()])
|
|
checkpoint = args.source / "model.safetensors"
|
|
with checkpoint.open("rb") as source:
|
|
digest = hashlib.file_digest(source, "sha256").hexdigest()
|
|
if digest != "87a4a970ce9aa79d5d92e71899ab034defcf13d93e5e4393ec0dc7db6d4ec048":
|
|
raise ValueError("Expected the official NVIDIA RE-USE checkpoint")
|
|
command = [str(args.audiocpp_gguf.resolve()), "--input", f"weights={checkpoint.resolve()}",
|
|
"--output", str(args.output.resolve()), "--type", args.type, "--family", "reuse",
|
|
"--model-spec", str(root / "model_specs/reuse.json"), "--root", str(args.source.resolve())]
|
|
if args.overwrite:
|
|
command.append("--overwrite")
|
|
with safe_open(checkpoint, framework="np") as source:
|
|
for name in source.keys():
|
|
shape = source.get_slice(name).get_shape()
|
|
if len(shape) == 1 or name.endswith(".A_log") or name.endswith(".conv1d.weight"):
|
|
command.extend(["--keep-type", f"weights/{name}=f32"])
|
|
logging.info("command=%s", command)
|
|
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
|
|
logging.info("%s", result.stdout)
|
|
result.check_returncode()
|
|
reader = gguf.GGUFReader(str(args.output))
|
|
with safe_open(checkpoint, framework="np") as source:
|
|
if {tensor.name for tensor in reader.tensors} != {"weights/" + key for key in source.keys()}:
|
|
raise ValueError("GGUF tensor inventory mismatch")
|
|
ranks = reader.get_field("audiocpp.tensor_ranks").contents()
|
|
shapes = reader.get_field("audiocpp.tensor_shapes").contents()
|
|
cursor = 0
|
|
for tensor, rank in zip(reader.tensors, ranks, strict=True):
|
|
original = source.get_tensor(tensor.name.removeprefix("weights/"))
|
|
if tuple(shapes[cursor:cursor + rank]) != original.shape:
|
|
raise ValueError(f"GGUF shape mismatch: {tensor.name}")
|
|
cursor += rank
|
|
if args.type == "f32" and tensor.tensor_type != gguf.GGMLQuantizationType.F32:
|
|
raise ValueError(f"GGUF type mismatch: {tensor.name}")
|
|
if tensor.tensor_type == gguf.GGMLQuantizationType.F32 and not np.array_equal(
|
|
tensor.data.reshape(-1).view(np.uint8), original.reshape(-1).view(np.uint8)):
|
|
raise ValueError(f"GGUF type/value mismatch: {tensor.name}")
|
|
if cursor != len(shapes):
|
|
raise ValueError("Trailing GGUF tensor shape metadata")
|
|
logging.info("PASS: %d tensors with matching shapes; retained F32 tensors byte-exact; sha256=%s",
|
|
len(reader.tensors), digest)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|