# /// script
# requires-python = ">=3.11,<3.12"
# dependencies = [
#   "torch==2.13.0", "torchvision==0.28.0", "onnx==1.22.0",
#   "numpy==2.4.6", "transformers==4.57.6", "timm==1.0.25",
#   "kornia==0.8.2", "einops==0.8.2", "safetensors==0.7.0",
# ]
# [tool.uv.sources]
# torch = { index = "pytorch" }
# torchvision = { index = "pytorch" }
# [[tool.uv.index]]
# name = "pytorch"
# url = "https://download.pytorch.org/whl/cu130"
# explicit = true
# ///
"""Export pinned official BiRefNet weights to fixed-resolution ONNX alpha."""

import argparse
import hashlib
import json
import sys
from pathlib import Path
from urllib.request import urlopen

import onnx
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForImageSegmentation

from onnx_weights import compress_weights, compress_exact_constants
from deform_sampling import deform_conv, validate_deform


SOURCES = {
    "lite": ("ZhengPeng7/BiRefNet_lite", "aa62cd87eafb9cc43056d08ef3615a14628b831d", 1024),
    "hr-matting": ("ZhengPeng7/BiRefNet_HR-matting", "5d6b6f8adcb5b417c871b1d84ceaae9871355b7f", 2048),
}
WEIGHT_HASHES = {
    "lite": "4417d89795250e698c3cb0ae8df15743810065f646f48a694fdfa7ca052d0815",
    "hr-matting": "a5a4de698739ea5e0e8bbab28e1b293dde95092b87a442d566cbc585c53cef55",
}
LICENSE_REVISION = "ebcc0bc8ec7fe919cec829f2dea656b3078acddc"


class AlphaModel(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, image):
        return self.model(image)[-1].sigmoid()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--variant", choices=SOURCES, required=True)
    parser.add_argument("--destination", type=Path, required=True)
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument("--fp32-output", type=Path)
    args = parser.parse_args()
    repository, revision, size = SOURCES[args.variant]
    checkpoint = args.checkpoint / "model.safetensors" if args.checkpoint else Path(
        hf_hub_download(repository, "model.safetensors", revision=revision)
    )
    with checkpoint.open("rb") as stream:
        if hashlib.file_digest(stream, "sha256").hexdigest() != WEIGHT_HASHES[args.variant]:
            raise RuntimeError("BiRefNet checkpoint SHA-256 mismatch")
    torch.set_num_threads(8)
    model = AutoModelForImageSegmentation.from_pretrained(
        str(args.checkpoint) if args.checkpoint else repository,
        revision=revision,
        trust_remote_code=True,
    ).eval()
    validate_deform()
    sys.modules[type(model).__module__].deform_conv2d = deform_conv
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AlphaModel(model).to(device)
    args.destination.mkdir(parents=True, exist_ok=True)
    output = args.destination / "model.onnx"
    with torch.inference_mode():
        torch.onnx.export(
            model, torch.zeros(1, 3, size, size, device=device),
            str(output),
            input_names=["image"], output_names=["alpha"],
            opset_version=17, dynamo=False, do_constant_folding=False,
        )
    graph = onnx.load(output)
    if args.fp32_output:
        onnx.save(graph, args.fp32_output)
    compress_weights(graph)
    compress_exact_constants(graph)
    onnx.checker.check_model(graph)
    onnx.save(graph, output)
    license_url = f"https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/{LICENSE_REVISION}/LICENSE"
    license_text = urlopen(license_url, timeout=60).read()
    if hashlib.sha256(license_text).hexdigest() != "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb":
        raise RuntimeError("BiRefNet license SHA-256 mismatch")
    (args.destination / "LICENSE.txt").write_bytes(license_text)
    with output.open("rb") as stream:
        checksum = hashlib.file_digest(stream, "sha256").hexdigest()
    record = {"repository": repository, "revision": revision, "input_size": size,
              "bytes": output.stat().st_size, "sha256": checksum,
              "checkpoint_sha256": WEIGHT_HASHES[args.variant]}
    print(json.dumps(record), flush=True)


if __name__ == "__main__":
    main()