Add the models as Core ML ML programs for anira's native Core ML engine, and the export script that makes them

7d4dc70 verified about 20 hours ago

2.7 kB

# SPDX-License-Identifier: Apache-2.0
# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = ["coremltools==9.0", "torch==2.7.0", "onnx2torch", "onnxsim", "onnx", "onnxruntime", "numpy"]
# ///
"""Exports the four hand tracking models to Core ML (.mlpackage, an ML program in float16 for the
Neural Engine, Apple platforms) with the ONNX models' tensor names, from the PyTorch modules
onnx_to_torch.py makes, and checks each against ONNX Runtime on the same random input, relative to
each output's peak. torch 2.7 is the newest coremltools 9.0 is tested with.
    uv run export_coreml.py <weights dir laid out like the HF repo>
writes <weights dir>/coreml/<model>.mlpackage.
"""

import os
import sys

import coremltools as ct
import numpy as np
import onnxruntime
import torch

from onnx_to_torch import MODELS, torch_module

TOLERANCE = 1e-2  # float16, relative to the output's peak


def export(weights, name):
    onnx_path = os.path.join(weights, "onnx", f"{name}.onnx")
    module, shapes = torch_module(onnx_path)
    generator = torch.Generator().manual_seed(0)
    inputs = tuple(torch.rand(*dims, generator=generator) for dims in shapes.values())
    session = onnxruntime.InferenceSession(onnx_path)
    input_names = [v.name for v in session.get_inputs()]
    output_names = [v.name for v in session.get_outputs()]

    exported = torch.export.export(module, inputs).run_decompositions({})
    model = ct.convert(exported, convert_to="mlprogram", compute_precision=ct.precision.FLOAT16,
                       compute_units=ct.ComputeUnit.ALL, minimum_deployment_target=ct.target.macOS13)
    spec = model.get_spec()
    for old, new in zip([i.name for i in spec.description.input], input_names):
        ct.utils.rename_feature(spec, old, new)
    for old, new in zip([o.name for o in spec.description.output], output_names):
        ct.utils.rename_feature(spec, old, new)
    model = ct.models.MLModel(spec, weights_dir=model.weights_dir)
    os.makedirs(os.path.join(weights, "coreml"), exist_ok=True)
    model.save(os.path.join(weights, "coreml", f"{name}.mlpackage"))

    expected = session.run(None, {n: x.numpy() for n, x in zip(input_names, inputs)})
    actual = model.predict({n: x.numpy() for n, x in zip(input_names, inputs)})
    worst = max(float(np.max(np.abs(e - actual[n]))) / (1.0 + float(np.max(np.abs(e))))
                for n, e in zip(output_names, expected))
    print(f"{name}: max |Core ML - ONNX Runtime| = {worst:.2e}")
    if worst > TOLERANCE:
        sys.exit(f"{name}: the export differs from ONNX Runtime by {worst:.2e}")


def main():
    for name in MODELS:
        export(sys.argv[1], name)


if __name__ == "__main__":
    main()

	# SPDX-License-Identifier: Apache-2.0
	# /// script
	# requires-python = ">=3.11,<3.13"
	# dependencies = ["coremltools==9.0", "torch==2.7.0", "onnx2torch", "onnxsim", "onnx", "onnxruntime", "numpy"]
	# ///
	"""Exports the four hand tracking models to Core ML (.mlpackage, an ML program in float16 for the
	Neural Engine, Apple platforms) with the ONNX models' tensor names, from the PyTorch modules
	onnx_to_torch.py makes, and checks each against ONNX Runtime on the same random input, relative to
	each output's peak. torch 2.7 is the newest coremltools 9.0 is tested with.

	uv run export_coreml.py <weights dir laid out like the HF repo>

	writes <weights dir>/coreml/<model>.mlpackage.
	"""

	import os
	import sys

	import coremltools as ct
	import numpy as np
	import onnxruntime
	import torch

	from onnx_to_torch import MODELS, torch_module

	TOLERANCE = 1e-2 # float16, relative to the output's peak


	def export(weights, name):
	onnx_path = os.path.join(weights, "onnx", f"{name}.onnx")
	module, shapes = torch_module(onnx_path)
	generator = torch.Generator().manual_seed(0)
	inputs = tuple(torch.rand(*dims, generator=generator) for dims in shapes.values())
	session = onnxruntime.InferenceSession(onnx_path)
	input_names = [v.name for v in session.get_inputs()]
	output_names = [v.name for v in session.get_outputs()]

	exported = torch.export.export(module, inputs).run_decompositions({})
	model = ct.convert(exported, convert_to="mlprogram", compute_precision=ct.precision.FLOAT16,
	compute_units=ct.ComputeUnit.ALL, minimum_deployment_target=ct.target.macOS13)
	spec = model.get_spec()
	for old, new in zip([i.name for i in spec.description.input], input_names):
	ct.utils.rename_feature(spec, old, new)
	for old, new in zip([o.name for o in spec.description.output], output_names):
	ct.utils.rename_feature(spec, old, new)
	model = ct.models.MLModel(spec, weights_dir=model.weights_dir)
	os.makedirs(os.path.join(weights, "coreml"), exist_ok=True)
	model.save(os.path.join(weights, "coreml", f"{name}.mlpackage"))

	expected = session.run(None, {n: x.numpy() for n, x in zip(input_names, inputs)})
	actual = model.predict({n: x.numpy() for n, x in zip(input_names, inputs)})
	worst = max(float(np.max(np.abs(e - actual[n]))) / (1.0 + float(np.max(np.abs(e))))
	for n, e in zip(output_names, expected))
	print(f"{name}: max \|Core ML - ONNX Runtime\| = {worst:.2e}")
	if worst > TOLERANCE:
	sys.exit(f"{name}: the export differs from ONNX Runtime by {worst:.2e}")


	def main():
	for name in MODELS:
	export(sys.argv[1], name)


	if __name__ == "__main__":
	main()