Ship MI-GAN's bare 512 generator as the panorama border filler

Sargsyan et al., ICCV 2023; MIT code and weights (models/LICENCE.md),
exported by tools/export-migan.sh at a fixed 1×4×512×512 from the
authors' checkpoint — six operator types, 28 MB, in LFS like the rest.
The package installs it beside the scene model and the APK unpacks it
with the others.
This commit is contained in:
2026-09-19 20:41:20 +02:00
parent 2fd7690b6f
commit 54a80e688c
6 changed files with 100 additions and 3 deletions
+65
View File
@@ -0,0 +1,65 @@
#!/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"