72 lines
3.2 KiB
Bash
Executable File
72 lines
3.2 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
|
||
#
|
||
# Since 2026-09-20 the file that ships is not this export but a fine-tune of
|
||
# it on panorama-border voids (docs/panorama.md §14), made in the
|
||
# `darkroom-infill` repository with `python -m infill.export`. This script
|
||
# still yields the stock generator — the fine-tune's starting point, and the
|
||
# model the page's "mirror depth" knob above zero was built around.
|
||
#
|
||
# 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"
|