# BiRefNet_lite-matting for onnxruntime-web's WebGPU EP: each deformable conv becomes one GridSample + MatMul (float math
# only), and the sampling grid's base is made at run time from a CumSum over the offsets (a data tensor), so constant
# folding can't bake a full-size grid per layer into the file (that made the fp16 file 150-177 MB).
import sys, torch, torch.nn.functional as F, onnx, onnxslim
from transformers import AutoModelForImageSegmentation
from onnxruntime.transformers.float16 import convert_float_to_float16

def deform_conv2d_rt(input, offset, weight, bias=None, stride=(1, 1), padding=(0, 0), dilation=(1, 1), mask=None):
    n, c, h, w = input.shape
    o, _, kh, kw = weight.shape
    sh, sw = (stride, stride) if isinstance(stride, int) else stride
    ph, pw = (padding, padding) if isinstance(padding, int) else padding
    dh, dw = (dilation, dilation) if isinstance(dilation, int) else dilation
    ho, wo = offset.shape[2], offset.shape[3]
    k = kh * kw
    ones = offset[:, :1] * 0 + 1                       # n, 1, ho, wo
    oy = torch.cumsum(ones, 2) - 1
    ox = torch.cumsum(ones, 3) - 1
    ky, kx = torch.meshgrid(torch.arange(kh, dtype=input.dtype), torch.arange(kw, dtype=input.dtype), indexing='ij')
    ky = (ky.reshape(1, k, 1, 1) * dh)
    kx = (kx.reshape(1, k, 1, 1) * dw)
    off = offset.reshape(n, k, 2, ho, wo)
    y = oy * sh - ph + ky + off[:, :, 0]
    x = ox * sw - pw + kx + off[:, :, 1]
    grid = torch.stack([x * (2 / (w - 1)) - 1, y * (2 / (h - 1)) - 1], -1).reshape(n, k * ho, wo, 2)
    s = F.grid_sample(input, grid, mode='bilinear', padding_mode='zeros', align_corners=True).reshape(n, c, k, ho, wo)
    if mask is not None:
        s = s * mask.reshape(n, 1, k, ho, wo)
    out = torch.matmul(weight.reshape(o, c * k), s.reshape(n, c * k, ho * wo)).reshape(n, o, ho, wo)
    if bias is not None:
        out = out + bias.reshape(1, o, 1, 1)
    return out

model = AutoModelForImageSegmentation.from_pretrained('models/birefnet', trust_remote_code=True).eval()
mod = sys.modules[type(model).__module__]
orig = mod.deform_conv2d
x = torch.randn(1, 8, 20, 24); off = torch.randn(1, 18, 20, 24) * 2; m = torch.rand(1, 9, 20, 24) * 2
wt = torch.randn(5, 8, 3, 3); b = torch.randn(5)
print('op diff', float((orig(x, off, wt, b, padding=(1, 1), mask=m) - deform_conv2d_rt(x, off, wt, b, padding=(1, 1), mask=m)).abs().max()))
ims = torch.randn(1, 3, 512, 512)
with torch.no_grad():
    a = model(ims)[-1]
    mod.deform_conv2d = deform_conv2d_rt
    b2 = model(ims)[-1]
print('model diff', float((a - b2).abs().max()))
for size in [int(s) for s in sys.argv[1:]] or [1024]:
    raw = f'models/birefnet-rt-{size}.onnx'
    with torch.no_grad():
        torch.onnx.export(model, torch.randn(1, 3, size, size), raw, opset_version=17,
                          input_names=['input_image'], output_names=['output_image'], dynamo=False)
    slim = onnxslim.slim(onnx.load(raw))
    m16 = convert_float_to_float16(slim, keep_io_types=True)
    out = f'models/birefnet-lite-matting-{size}-rt-fp16.onnx'
    onnx.save(m16, out)
    ops = {}
    for nd in m16.graph.node: ops[nd.op_type] = ops.get(nd.op_type, 0) + 1
    print('saved', out, len(m16.graph.node), 'nodes', sorted(ops.items(), key=lambda t: -t[1])[:12], flush=True)