# Shrinks an exported BiRefNet: BiRefNet runs its backbone twice (full and half scale), and the export writes the backbone's
# weights once per call, some of them transposed. Identical initializers are merged, and a transposed copy becomes a
# Transpose of the original (ORT folds it back when the session is created, so it costs nothing per frame). Node names
# are made unique (the fp16 conversion repeats a few).
import sys, hashlib, onnx, numpy as np
from onnx import helper, numpy_helper

def digest(a):
    return hashlib.sha1(np.ascontiguousarray(a).tobytes()).hexdigest()

def toposort(nodes, available):
    """The nodes in an order where each comes after what makes its inputs (the fp16 conversion appends its casts)."""
    have, out, pending = set(available) | {''}, [], list(nodes)
    while pending:
        rest = []
        for nd in pending:
            if all(i in have for i in nd.input):
                out.append(nd)
                have.update(nd.output)
            else:
                rest.append(nd)
        if len(rest) == len(pending): raise ValueError(f'{len(rest)} nodes have inputs nothing makes, e.g. {rest[0].name}')
        pending = rest
    return out

def shrink(m):
    seen, rename, keep, transposed = {}, {}, [], []
    for t in m.graph.initializer:
        a = numpy_helper.to_array(t)
        key = (t.data_type, a.shape, digest(a))
        if key in seen:
            rename[t.name] = seen[key]
            continue
        if a.ndim == 2 and a.size >= 4096:
            kt = (t.data_type, a.T.shape, digest(a.T))
            if kt in seen:
                transposed.append(helper.make_node('Transpose', [seen[kt]], [t.name], perm=[1, 0], name=f'{t.name}_from_transpose'))
                continue
        seen[key] = t.name
        keep.append(t)
    del m.graph.initializer[:]
    m.graph.initializer.extend(keep)
    for nd in m.graph.node:
        for i, n in enumerate(nd.input):
            if n in rename: nd.input[i] = rename[n]
    # The fp16 conversion casts a value once per op it keeps in fp32, under one name: identical casts collapse to the first.
    made, unique = {}, []
    for nd in list(m.graph.node):
        prev = made.get(nd.output[0]) if nd.output else None
        if prev is not None:
            if prev.op_type != nd.op_type or list(prev.input) != list(nd.input): raise ValueError(f'two different nodes make {nd.output[0]}')
            continue
        for o in nd.output: made[o] = nd
        unique.append(nd)
    del m.graph.node[:]
    m.graph.node.extend(unique)
    nodes = toposort(transposed + list(m.graph.node), {i.name for i in m.graph.input} | {t.name for t in m.graph.initializer})
    names = set()
    for i, nd in enumerate(nodes):
        if not nd.name or nd.name in names: nd.name = f'{nd.name or nd.op_type}_{i}'
        names.add(nd.name)
    del m.graph.node[:]
    m.graph.node.extend(nodes)
    return m, len(rename), len(transposed)

if __name__ == '__main__':
    m, n, t = shrink(onnx.load(sys.argv[1]))
    onnx.checker.check_model(m)
    onnx.save(m, sys.argv[2])
    print('merged', n, 'transposed', t)