feat(studio): HEAL invents what was under the brush with MI-GAN, and keeps the blurred fill for the machines that cannot run it

The dust brush filled its hole with a blur of the hole's own rim — the light
beside it, carried in. Right for a speck on a wall, and visibly thin over
anything with structure: a stroke across a shirt came back as the average of a
shirt. The fill is now a learned one.

engine/inpaint.ts runs MI-GAN (andraniksargsyan/migan, the same
`migan_pipeline_v2.onnx` inpaint-web ships) through onnxruntime-web on the
visitor's own machine: 27MB of weights in public/models, WebGPU where the graph
fits the provider, the processor otherwise, one warm-up pass so the first repair
is not the one that pays for the kernels. It is handed the masked crop and
nothing else — the model's mask polarity is 255 KEEP / 0 HOLE, the opposite of
the obvious guess, and the other way round repaints the frame and leaves the
hole (measured).

Step 6e of the renderer cuts the crop at the resolution the model will actually
run at (512 on its longest side — the graph is fully convolutional and would
take a whole frame, and ten seconds over one), builds the mask from the brush's
own circles, and draws the answer into the fill surface at its own place in the
box. The composite, the feathered mask and the untouched frame around it are
unchanged, so nothing outside the mask moves — and when the model cannot be had
the blurred fill is still built first and still stands. The preview gets this
for free: renderPhoto was already async and the preview was already debounced.

Node probe: the model + polarity + read/write path run; heal-check.mjs keeps
measuring the fallback and now asserts the wiring and the polarity. Browser:
engine/inpaint.ts in a real Chromium returns a patch (hole moves, kept pixels
within 0.02 of a level), and the built app renders a 42-spot stroke in 4.7s on
single-threaded wasm with no fallback warning — 1.15s and the warning when the
model is taken away.

docker/frontend/src/engine/inpaint.ts

@@ -0,0 +1,239 @@
+import type * as Ort from 'onnxruntime-web';
+
+// FX > HEAL, the half that invents. The dust brush used to be filled with a
+// blur of the photo's own rim — the light next to the hole, carried in, which is
+// right for a speck of dust on a wall and thin over anything with structure: a
+// hole across a shirt comes back as the average of a shirt.
+//
+// The fill is now a learned one. MI-GAN (andraniksargsyan/migan, the
+// `migan_pipeline_v2.onnx` weights inpaint-web ships) is 27MB and runs through
+// onnxruntime-web on the visitor's own machine, like the upscaler next door and
+// for the same reason: the photo never leaves the device. It is only asked for
+// what a GAN is actually good at — the pixels under the brush, and not one
+// pixel more: everything outside the mask is the frame as rendered, put back by
+// the composite in the renderer (see step 6e). A model that is wrong about a
+// corner of the crop costs nothing, because that corner is never read.
+//
+// The runtime and the model are static files beside the app; nothing here
+// changes what the server can do.
+//
+// The mask polarity is MI-GAN's own and is not the intuitive one: 255 is KEEP
+// and 0 is the HOLE to invent — the model repaints exactly where the mask is
+// zero. Measured on the stroke probe (tmp-migan-probe.mjs, since deleted):
+// mean |before-after| over the hole 43.12, over everything the mask kept 0.01.
+// The other way round the frame comes back repainted and the hole untouched,
+// which is the collage this file exists to avoid.
+
+const MODEL_URL = '/models/migan_pipeline_v2.onnx';
+// The wasm dir, not the file: the runtime picks its own name inside it.
+const WASM_DIR = '/wasm/ort/';
+// What a HEAL toggle warms: the runtime's own files and the model. Same list as
+// the upscaler's, whose runtime this is — the bytes are shared.
+const RUNTIME_FILES = ['ort-wasm-simd-threaded.jsep.mjs', 'ort-wasm-simd-threaded.jsep.wasm'];
+// The longest side the model is ever handed, before the renderer's crop is
+// scaled down to it. The graph is fully convolutional and would take a whole
+// frame; it would also take ten seconds over one, per keystroke of the brush.
+// 512 is the probe's own size and about a second on four threads of wasm, and a
+// repair is a few brush widths across — the crop has all the context it needs
+// long before this. `inpaintCrop` scales what it is given to fit; the renderer
+// draws the answer back to the crop's own rectangle.
+export const INPAINT_MAX_SIDE = 512;
+// The one throwaway run when the session opens: 64 pixels square, but a whole
+// number of the graph's own stride, so the kernels it builds are the kernels the
+// repairs use.
+const WARM_SIDE = 64;
+
+interface Loaded {
+  ort: typeof Ort;
+  session: Ort.InferenceSession;
+}
+
+let loaded: Promise<Loaded> | null = null;
+
+// One warm-up run through the graph, answering whether the session it was given
+// can actually execute the model. WebGPU is worth having — a repair lands in
+// well under a second instead of a second and a half — but a provider the graph
+// does not fit is a session that opens happily and throws on the first pixel, so
+// the trial is a real run and not a `create` that returned.
+async function warmsUp(ort: typeof Ort, session: Ort.InferenceSession): Promise<boolean> {
+  try {
+    await session.run({
+      [session.inputNames[0]]: new ort.Tensor('uint8', new Uint8Array(3 * WARM_SIDE * WARM_SIDE), [
+        1, 3, WARM_SIDE, WARM_SIDE,
+      ]),
+      [session.inputNames[1]]: new ort.Tensor('uint8', new Uint8Array(WARM_SIDE * WARM_SIDE).fill(255), [
+        1, 1, WARM_SIDE, WARM_SIDE,
+      ]),
+    });
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+// WebGPU first, the processor as the answer to every way that can fail: no
+// adapter, no `shader-f16` needed here but an op the provider never
+// implemented, a driver that opens a session and dies on the first dispatch.
+async function openSession(ort: typeof Ort): Promise<Ort.InferenceSession> {
+  try {
+    const gpu = await ort.InferenceSession.create(MODEL_URL, { executionProviders: ['webgpu'] });
+    if (await warmsUp(ort, gpu)) return gpu;
+    await gpu.release();
+  } catch {
+    // Whatever went wrong, the wasm provider is the fallback and not a failure.
+  }
+  return ort.InferenceSession.create(MODEL_URL, { executionProviders: ['wasm'] });
+}
+
+// Loads the runtime and the model once per page, and runs one warm-up pass so
+// the first repair is not the one that pays for the kernels.
+function load(): Promise<Loaded> {
+  if (!loaded) {
+    loaded = (async () => {
+      const ort = await import('onnxruntime-web');
+      ort.env.wasm.wasmPaths = WASM_DIR;
+      // Same rule as the upscaler's: the threads need COOP+COEP for their
+      // SharedArrayBuffer, and asking for more than one without isolation is a
+      // throw rather than a fallback. Four, because at 512 pixels square the
+      // graph is done before a wider pool has paid for its own start-up.
+      ort.env.wasm.numThreads = self.crossOriginIsolated
+        ? Math.min(4, navigator.hardwareConcurrency || 1)
+        : 1;
+      ort.env.webgpu.powerPreference = 'high-performance';
+      const session = await openSession(ort);
+      return { ort, session };
+    })().catch((err) => {
+      // A failed load must not poison the next attempt.
+      loaded = null;
+      throw err;
+    });
+  }
+  return loaded;
+}
+
+// Start fetching the model before the brush that needs it, so the visitor waits
+// for the repair and not for the download. Called when HEAL is switched on.
+//
+// The bytes, and only the bytes: importing the runtime starts its thread pool —
+// a Web Worker each, a copy of the 27MB binary each — and a pool that then
+// belongs to the page until it dies, whether or not a stroke is ever painted.
+// Switching a tool on is not the work; the download is, and the browser's cache
+// is what keeps it.
+export function preloadInpaint(): void {
+  for (const file of [...RUNTIME_FILES.map((f) => WASM_DIR + f), MODEL_URL]) {
+    void fetch(file).catch(() => {});
+  }
+}
+
+// The page is going away, or coming back from the back/forward cache: hand the
+// session back rather than leaving the browser to reap its buffers and its
+// thread pool with the page. `loaded` goes with it, so a page that does come
+// back loads the runtime again instead of calling into a released session — the
+// files are cached, so that is a second.
+export async function releaseInpaint(): Promise<void> {
+  const pending = loaded;
+  loaded = null;
+  if (!pending) return;
+  try {
+    await (await pending).session.release();
+  } catch {
+    // A load that failed, or a runtime too old to have `release`.
+  }
+}
+
+if (typeof addEventListener === 'function') {
+  addEventListener('pagehide', () => {
+    void releaseInpaint();
+  });
+}
+
+// The pixels under the brush, invented from the pixels around them.
+//
+// `rgba` is the crop the renderer cut out of the frame, at its own resolution;
+// `hole` is one byte per pixel, non-zero for the pixels the brush painted. The
+// answer is the same crop in RGBA at the same size, or null when the model could
+// not be had — the caller keeps its blur fill in that case, and a repair that is
+// merely soft beats no repair at all.
+//
+// A crop larger than INPAINT_MAX_SIDE on its longest side is scaled down to fit
+// and the answer is scaled back up by the caller. A brush stroke is a soft shape
+// made of soft shapes; what a factor of two of resample costs it is nothing
+// beside the seconds it saves.
+export async function inpaintCrop(
+  rgba: Uint8Array | Uint8ClampedArray,
+  width: number,
+  height: number,
+  hole: Uint8Array
+): Promise<Uint8Array | null> {
+  let w = Math.max(8, Math.round(width));
+  let h = Math.max(8, Math.round(height));
+  const scale = Math.min(1, INPAINT_MAX_SIDE / Math.max(w, h));
+  // Whole multiples of the graph's own stride, so the last block of the encoder
+  // is a block and not a remainder: rounded down, since the renderer's crop is
+  // wider than the mask by its own margin and losing a sliver of that margin
+  // costs the mask nothing.
+  w = Math.max(8, Math.floor((w * scale) / 8) * 8);
+  h = Math.max(8, Math.floor((h * scale) / 8) * 8);
+
+  // NCHW, 0..255 RGB — what the model was trained to read — and the mask in its
+  // own polarity: 255 everywhere the frame is known, 0 over the hole.
+  const src = new Uint8Array(3 * w * h);
+  const mask = new Uint8Array(w * h).fill(255);
+  const sx = width / w;
+  const sy = height / h;
+  const plane = w * h;
+  for (let y = 0; y < h; y++) {
+    // Nearest source pixel rather than a box filter: the crop is the frame at
+    // its own resolution whenever it fits, so the two agree pixel for pixel and
+    // this loop never runs at all in the common case.
+    const syy = Math.min(height - 1, Math.floor(y * sy));
+    for (let x = 0; x < w; x++) {
+      const sxx = Math.min(width - 1, Math.floor(x * sx));
+      const s = (syy * width + sxx) * 4;
+      const d = y * w + x;
+      src[d] = rgba[s];
+      src[plane + d] = rgba[s + 1];
+      src[2 * plane + d] = rgba[s + 2];
+      if (hole[syy * width + sxx]) mask[d] = 0;
+    }
+  }
+
+  try {
+    const { ort, session } = await load();
+    const out = await session.run({
+      [session.inputNames[0]]: new ort.Tensor('uint8', src, [1, 3, h, w]),
+      [session.inputNames[1]]: new ort.Tensor('uint8', mask, [1, 1, h, w]),
+    });
+    const tensor = out[session.outputNames[0]];
+    const oplane = w * h;
+    const data = tensor.data as Uint8Array;
+    // Back to RGBA, opaque: the crop covers the frame, which has no transparency
+    // of its own to carry.
+    const patch = new Uint8Array(4 * w * h);
+    for (let i = 0; i < oplane; i++) {
+      patch[i * 4] = data[i];
+      patch[i * 4 + 1] = data[oplane + i];
+      patch[i * 4 + 2] = data[2 * oplane + i];
+      patch[i * 4 + 3] = 255;
+    }
+    return patch;
+  } catch (err) {
+    // A model that cannot run is a soft repair, not a broken one — and this one
+    // line is what keeps the fallback diagnosable instead of invisible, because
+    // the picture that comes out of it looks like a workable repair.
+    console.warn('[inpaint] model unavailable, keeping the blurred fill', err);
+    return null;
+  }
+}
+
+// The size `inpaintCrop` will actually run at, for a caller that has to read the
+// crop back at the model's resolution rather than at the frame's. The renderer
+// uses it to cut the crop it hands over: a 3000px box read back at 3000px and
+// scaled here would be 27MB of pixels moved for a repair that ends up 512 wide.
+export function inpaintCropSize(width: number, height: number): { w: number; h: number } {
+  const scale = Math.min(1, INPAINT_MAX_SIDE / Math.max(1, Math.max(width, height)));
+  return {
+    w: Math.max(8, Math.floor((Math.max(1, Math.round(width)) * scale) / 8) * 8),
+    h: Math.max(8, Math.floor((Math.max(1, Math.round(height)) * scale) / 8) * 8),
+  };
+}