#!/usr/bin/env python3
"""Widen SAM 3D Body's decoder from bfloat16 to float32.

ONNX Runtime's CPU provider has no MatMul kernel for bfloat16, so the decoder as
shipped loads and then refuses at its first matrix multiply -- which is what
`peasanttide_portrait` hit and is the only reason this script exists. Every
bfloat16 is exactly representable as a float32 (it *is* a float32 with the low
sixteen bits of the mantissa cleared), so this is a widening and not a
conversion: no retraining, no quantisation, and the same numbers to the bit.

Three things have to move together or the graph will not type-check: the
initialisers, the tensors that live inside node attributes, and every `Cast`
whose target is bfloat16. The stale `value_info` is dropped and re-inferred
rather than edited, because editing it by hand is how a graph ends up half
converted.

    uv run --with onnx --with numpy python tools/widen_decoder.py in.onnx out.onnx
"""

import sys

import numpy as np
import onnx
from onnx import TensorProto, numpy_helper


def widen(tensor: TensorProto) -> bool:
    """Rewrites one tensor in place. Answers whether it was a bfloat16 one."""
    if tensor.data_type != TensorProto.BFLOAT16:
        return False
    if tensor.raw_data:
        raw = np.frombuffer(tensor.raw_data, dtype=np.uint16)
    else:
        raw = np.array(tensor.int32_data, dtype=np.uint32).astype(np.uint16)
    wide = (raw.astype(np.uint32) << 16).view(np.float32)
    shape = [dim for dim in tensor.dims]
    tensor.CopyFrom(numpy_helper.from_array(wide.reshape(shape), tensor.name))
    return True


def main() -> None:
    source, into = sys.argv[1], sys.argv[2]
    model = onnx.load(source)

    initialisers = sum(widen(one) for one in model.graph.initializer)
    attributes = 0
    casts = 0
    for node in model.graph.node:
        for attribute in node.attribute:
            if attribute.type == onnx.AttributeProto.TENSOR:
                attributes += widen(attribute.t)
            elif attribute.type == onnx.AttributeProto.TENSORS:
                for one in attribute.tensors:
                    attributes += widen(one)
            elif attribute.name == "to" and attribute.i == TensorProto.BFLOAT16:
                attribute.i = TensorProto.FLOAT
                casts += 1
    for value in list(model.graph.input) + list(model.graph.output):
        if value.type.tensor_type.elem_type == TensorProto.BFLOAT16:
            value.type.tensor_type.elem_type = TensorProto.FLOAT
    del model.graph.value_info[:]

    model = onnx.shape_inference.infer_shapes(model, strict_mode=False)
    onnx.save(model, into)
    print(
        f"{initialisers} initialisers, {attributes} attribute tensors and "
        f"{casts} casts widened -> {into}"
    )


if __name__ == "__main__":
    main()