Files
DarkRoom/tools/export-migan.sh
T
dtourolle 84fade99ec Put the developer docs under docs/dev and index the folder for users first
docs/ had 26 developer documents flat beside the manual, and the two
audiences are very differently sized: most readers want the manual and
the gesture reference, a few want the register, the designs and the
measurements. The manual and gestures.md stay at the top; everything for
someone changing the code moves to docs/dev/, and the two documents that
name their own successors — the v0.1 milestone and the UI-refinement plan
— go to docs/dev/archive/ rather than being deleted, since both are still
cited. docs/README.md is the index, users first.

Every reference follows: code comments, Cargo manifests, the workflows,
the pre-commit hook, the bench and traceability tools (which locate the
repo root by docs/dev/requirements.md now), packaging, the Docker READMEs,
CLAUDE.md, CONTRIBUTING.md and the README. The matrix links one level
deeper and is regenerated. Links out of the moved documents into the tree
gain a level; a link checker over every Markdown file finds none broken.
2026-09-20 21:16:03 +02:00

72 lines
3.3 KiB
Bash
Executable File
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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/dev/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"