# Converts SigLIP 2's vision encoder to a Core ML model for modules/Labels on macOS, then checks it
# against the ONNX model mcap runs elsewhere, on a folder of images. Run once, on a Mac; mcap only uses
# the .mlpackage it writes (into modules/Labels/models, where the build copies it from).
#
#   python3 -m venv ~/mcap-coreml
#   ~/mcap-coreml/bin/pip install torch==2.7.0 transformers coremltools==9.0 onnxruntime pillow sentencepiece protobuf
#   ~/mcap-coreml/bin/python tools/export_siglip_coreml.py --check FOLDER_OF_IMAGES
#
# torch 2.7.0 is the newest coremltools 9.0 was tested with. The model is the vision tower of
# google/siglip2-base-patch16-384, from pixel_values [1, 3, 384, 384] (RGB resized to 384 x 384,
# scaled to -1..1) to pooler_output [1, 768], as the ONNX one; an ML Program in float16, for the GPU.
# Measured on an M5 Pro: 6.7 ms an image on the GPU, against 67 ms for onnxruntime on the CPU, with
# vectors within cosine 0.99994 of the CPU's. The Neural Engine gets float16 badly wrong for this
# model (cosine 0.72), so mcap runs it on the CPU and GPU only.

import argparse, glob, os, sys, time
import numpy as np
import torch
from PIL import Image
import coremltools as ct
from transformers import AutoModel
from transformers.utils import logging

logging.set_verbosity_error()

CHECKPOINT = "google/siglip2-base-patch16-384"
SIZE = 384
HERE = os.path.dirname(os.path.abspath(__file__))
MODELS = os.path.join(HERE, "..", "modules", "Labels", "models")
OUTPUT = os.path.join(MODELS, "siglip2-base-patch16-384-vision.mlpackage")
ONNX = os.path.join(MODELS, "siglip2-base-patch16-384-vision.onnx")


# coremltools 9.0 casts a folded one-element tensor with int(array), which NumPy 2 refuses ("only
# 0-dimensional arrays can be converted"); the same cast, with the element taken out first.
from coremltools.converters.mil.frontend.torch import ops as ct_ops
_original_cast = ct_ops._cast
def _cast_one_element(context, node, dtype, dtype_name):
    x = ct_ops._get_inputs(context, node, expected=1)[0]
    if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
        context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
        return
    _original_cast(context, node, dtype, dtype_name)
ct_ops._cast = _cast_one_element


class Pooled(torch.nn.Module):
    """The vision tower, from pixel values to the pooled vector, as the ONNX graph has it."""

    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, pixel_values):
        return self.model(pixel_values=pixel_values).pooler_output


def load_images(folder, limit):
    paths = sorted(p for p in glob.glob(os.path.join(folder, "**", "*"), recursive=True)
                   if p.lower().endswith((".png", ".jpg", ".jpeg", ".webp")))[:limit]
    images = []
    for p in paths:
        image = Image.open(p).convert("RGB").resize((SIZE, SIZE), Image.BICUBIC)
        a = np.asarray(image, dtype=np.float32) / 255.0
        images.append(((a - 0.5) / 0.5).transpose(2, 0, 1))
    return np.stack(images).astype(np.float32)


def worst_cosine(a, b):
    a = a / np.linalg.norm(a, axis=-1, keepdims=True)
    b = b / np.linalg.norm(b, axis=-1, keepdims=True)
    return float(np.min(np.sum(a * b, axis=-1)))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--check", metavar="FOLDER", help="images to check the Core ML model against the ONNX one on")
    parser.add_argument("--limit", type=int, default=32)
    args = parser.parse_args()

    model = Pooled(AutoModel.from_pretrained(CHECKPOINT, torch_dtype=torch.float32).vision_model.eval()).eval()
    example = torch.zeros(1, 3, SIZE, SIZE)
    with torch.no_grad():
        traced = torch.jit.trace(model, example)
    start = time.perf_counter()
    ml = ct.convert(traced, inputs=[ct.TensorType(name="pixel_values", shape=(1, 3, SIZE, SIZE))],
                    outputs=[ct.TensorType(name="pooler_output")],
                    convert_to="mlprogram", compute_precision=ct.precision.FLOAT16,
                    minimum_deployment_target=ct.target.macOS14)
    ml.short_description = f"{CHECKPOINT} vision encoder (pooler_output), float16, for mcap"
    ml.save(OUTPUT)
    print(f"Wrote {OUTPUT} in {time.perf_counter() - start:.1f} s.")

    if not args.check:
        return
    import onnxruntime as ort
    images = load_images(args.check, args.limit)
    if not len(images):
        sys.exit(f"No images in {args.check}.")
    session = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"])
    name = session.get_inputs()[0].name
    reference = np.concatenate([session.run(["pooler_output"], {name: images[i:i + 1]})[0] for i in range(len(images))])
    coreml = ct.models.MLModel(OUTPUT, compute_units=ct.ComputeUnit.CPU_AND_GPU)
    coreml.predict({"pixel_values": images[:1]})
    start = time.perf_counter()
    got = np.concatenate([coreml.predict({"pixel_values": images[i:i + 1]})["pooler_output"] for i in range(len(images))])
    per = (time.perf_counter() - start) / len(images)
    cosine = worst_cosine(got.astype(np.float32), reference)
    print(f"{len(images)} images: {per * 1000:.1f} ms an image on the GPU, worst cosine to the ONNX model on the CPU {cosine:.5f}.")
    if cosine < 0.999:
        sys.exit("The Core ML model is too far from the ONNX one.")


if __name__ == "__main__":
    main()