"""Halve faceage_dino_fp32.onnx to fp16 with fp32 I/O.

    python to_fp16.py                       # plain conversion
    python to_fp16.py LayerNormalization Softmax ...   # with an op_block_list
"""

import hashlib
import os
import sys

import onnx
from onnxconverter_common import float16

SRC = r"E:/projects/faceage-onnx/source/faceage_dino_fp32.onnx"
DST = r"E:/projects/faceage-onnx/faceage-dino-fp16.onnx"


def sha256(path):
    d = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            d.update(chunk)
    return d.hexdigest()


def main():
    extra = sys.argv[1:]
    model = onnx.load(SRC)
    kwargs = {"keep_io_types": True}
    if extra:
        kwargs["op_block_list"] = list(float16.DEFAULT_OP_BLOCK_LIST) + extra
    print("extra ops kept in fp32:", extra)
    converted = float16.convert_float_to_float16(model, **kwargs)
    onnx.save(converted, DST, save_as_external_data=False)
    extra = [
        f for f in os.listdir(os.path.dirname(DST))
        if not f.endswith(".onnx") and not os.path.isdir(os.path.join(os.path.dirname(DST), f))
    ]
    print("non-.onnx files next to the output (external data would show here):", extra)
    print("fp32", SRC, os.path.getsize(SRC), sha256(SRC))
    print("fp16", DST, os.path.getsize(DST), sha256(DST))


if __name__ == "__main__":
    main()