#!/usr/bin/env bash
# Re-export the border-fill model that ships in models/inpaint/.
#
# The .onnx is committed (D14, FR-MRG-4), so this is not part of any build —
# it exists so the committed artefact is reproducible rather than a binary
# someone once produced and nobody can regenerate. Run it when bumping the
# model.
#
#   ./tools/export-migan.sh    # -> models/inpaint/migan-512.onnx
#
# Requires `uv`. Everything else is fetched into a throwaway venv, including
# a CPU-only torch and `gdown` for the checkpoint, which the authors keep on
# Google Drive (models/LICENCE.md has the licence; it is MIT).
#
# ## What is exported, and what is not
#
# The *bare* 512 generator, at a fixed 1×4×512×512: channel 0 is the mask
# minus a half (1 where the picture is known), channels 1–3 the RGB in −1..1
# with the unknown pixels zeroed; out come three planes in −1..1. The
# authors' "pipeline" ONNX — crop around the mask, resize, blend, all in the
# graph with dynamic shapes — is what tract cannot load, and every one of
# those steps is done in Rust (`dr_pano::fill`), where a tile of a panorama
# border needs different context from a brush stroke anyway.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO="$(cd "${HERE}/.." && pwd)"
OUT="${REPO}/models/inpaint"

# Not `mktemp -d` under /tmp: a tmpfs, and torch is a gigabyte.
WORK="$(mktemp -d -p "${TMPDIR:-/var/tmp}")"
trap 'rm -rf "${WORK}"' EXIT

echo "==> exporting MI-GAN in ${WORK}"
cd "${WORK}"
git clone -q --depth 1 https://github.com/Picsart-AI-Research/MI-GAN.git migan
uv venv -q --python 3.12 venv
VIRTUAL_ENV="${WORK}/venv" uv pip install -q --index-url https://download.pytorch.org/whl/cpu torch
VIRTUAL_ENV="${WORK}/venv" uv pip install -q onnx onnxslim gdown

# The 512 Places2 checkpoint, from the authors' Drive folder.
"${WORK}/venv/bin/gdown" --quiet "https://drive.google.com/uc?id=1D_YCuCgo20S2256sqpedsmENNm2WMtVY" -O "${WORK}/migan_512_places2.pt"

VIRTUAL_ENV="${WORK}/venv" "${WORK}/venv/bin/python" - "${WORK}/migan" "${WORK}/migan_512_places2.pt" "${WORK}/migan-512.onnx" <<'PY'
import sys, torch, onnx, onnxslim
sys.path.insert(0, sys.argv[1])
from lib.model_zoo.migan_inference import Generator as MIGAN

repo, ckpt, out = sys.argv[1], sys.argv[2], sys.argv[3]
model = MIGAN(resolution=512).eval()
model.load_state_dict(torch.load(ckpt, map_location="cpu"))
torch.onnx.export(model, torch.zeros(1, 4, 512, 512), out, opset_version=17, dynamo=False,
                  input_names=["input"], output_names=["output"],
                  dynamic_axes=None, do_constant_folding=True)
m = onnxslim.slim(onnx.load(out))
onnx.checker.check_model(m)
onnx.save(m, out)
print("ops:", sorted({n.op_type for n in m.graph.node}))
for o in m.graph.output:
    print("out", o.name, [d.dim_value for d in o.type.tensor_type.shape.dim])
PY

mkdir -p "${OUT}"
cp "${WORK}/migan-512.onnx" "${OUT}/migan-512.onnx"
echo "==> ${OUT}/migan-512.onnx"