dtourolle
54a80e688c
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.
66 lines
2.9 KiB
Bash
Executable File
66 lines
2.9 KiB
Bash
Executable File
#!/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"
|