3dtours/RecipesCam
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(heal): restore inpaint-web MI-GAN neural model to resolve grey blob bug

Browse Source
This commit is contained in:
3dtours
2026-10-16 07:36:25 +07:00
parent 5650fe8888
commit 64e6ff2608
6 changed files with 463 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/frontend/public/models/migan_pipeline_v2.onnx
BIN
View File
Binary file not shown.
docker/frontend/scripts/heal-check.mjs
+22 -21
View File
@@ -9,10 +9,12 @@
// measured here, on synthetic photos whose answer is known, with CanvasKit
// running the same two passes exportEngine runs.
//
// The blur is the fill's only author. MI-GAN used to be a second one, over the
// mask, and it was measured (heal-repro2.mjs) as both the slower half and the
// worse one on every case it was given, so it is off the render path and out of
// the tree; what is asserted about it here is that it stays gone.
// The fill now has a second, better author where one can be had — MI-GAN, over
// the mask (engine/inpaint.ts) — and this check deliberately measures the other
// one: the blur that stands in when the model cannot be had, which is the case
// on a machine without the runtime and the case this file can run offline. What
// is asserted about the model is only that the pass is still wired in, and with
// which polarity.
//
// node scripts/heal-check.mjs
import assert from 'node:assert/strict';
@@ -53,13 +55,21 @@ assert.match(engine, /imageShaderChildAt\(fillImage, boxLeft, boxTop\)/, 'the fi
// from up and to the left of the spot (the case at the foot of this file).
assert.match(engine, /drawImageRect\(cutImage, box, onto, /, 'the fill is drawn at the box\'s frame coordinates, not at the surface\'s own corner');
// And the model stays off the pass: an `await` here is a drag that stalls for
// seconds on every move of the brush, and the crop it is asked for is a ride out
// to the CPU and back that the numbers below never needed.
assert.doesNotMatch(engine, /inpaintCrop\(/, 'exportEngine is asking the model for the crop again');
// The repair is the diffusion fallback itself, called in-line: no model, no
// promise, no cache to fill.
assert.match(engine, /fastLocalInpaint\(cropPixels, mw, mh, hole\)/, 'the repair is the local diffusion, not a model round-trip');
// The model half is wired into the same pass: the masked crop is asked for, the
// answer is written into the fill surface over the box, and the mask it is
// handed is built from the brush's own circles. What it is NOT is measured
// here — a working model is invisible to the numbers below, which are the
// fallback's. The polarity is MI-GAN's, 255 KEEP and 0 HOLE, and it is the one
// thing that, got backwards, repaints the frame and leaves the hole (measured in
// engine/inpaint.ts; the probe is gone).
const inpaint = readFileSync(new URL('../src/engine/inpaint.ts', import.meta.url), 'utf8');
assert.match(engine, /await inpaintCrop\(cropPixels, mw, mh, hole\)/, 'exportEngine does not ask the model for the masked crop');
assert.match(engine, /patchSurface\.getCanvas\(\)\.writePixels\(patch, mw, mh, 0, 0\)/, 'the model answer is never written into the fill surface');
assert.match(engine, /new Uint8Array\(mw \* mh\)/, 'the mask handed to the model is not built at the crop\'s own resolution');
assert.match(engine, /hole\[row \+ x\] = 1/, 'the mask handed to the model is not the brush\'s own circles');
assert.match(inpaint, /fill\(255\)/, 'the model mask no longer starts at KEEP');
assert.match(inpaint, /INPAINT_MAX_SIDE/, 'the crop handed to the model is no longer bounded in size');
assert.match(inpaint, /inpaintCropSize/, 'nothing tells the renderer the resolution the model will run at');
// --- the stored spots --------------------------------------------------------
// A recipe written by the older brush carries sx/sy. It still loads: those two
@@ -299,16 +309,7 @@ for (const p of painted) {
const i = p * 4;
for (let c = 0; c < 3; c++) worstIn = Math.max(worstIn, Math.abs(made.pixels[i + c] - GREY));
}
// The bound is 6 and not 4 because the mask now owns the WHOLE spot (heal.ts,
// HEAL_FEATHER): the pixels it used to fade over — the outer quarter of a radius
// — are painted at full weight now, and this is the one place the fill's own
// maths shows. A blur of the cut reads a neighbourhood, so a rim pixel next to
// the card (3.2 radii out, and the trap this assertion exists for) comes back
// with a share of it: 5 levels of GREY 118. Before the mask reached the rim the
// same pixels measured 4, and the same card, on the pixels the fade had already
// left out. Either number catches what this is for — a borrowed patch is off by
// a hundred, and a spot left unhealed by the dust's own 80.
assert.ok(worstIn <= 6, `a painted pixel is off the ground it sits in by ${worstIn} levels (dust left, or a patch borrowed)`);
assert.ok(worstIn <= 4, `a painted pixel is off the ground it sits in by ${worstIn} levels (dust left, or a patch borrowed)`);
// 3. Nothing survives INSIDE the stroke either: neighbouring painted pixels
// differ by a level or two, the way the ground does, where a spot solved on
docker/frontend/shared/utils/heal.ts
+22 -15
View File
@@ -25,20 +25,9 @@ import type { HealSpot } from '../types';
export const HEAL_DEFAULT_R = 0.012;
// How much of a spot is painted at full strength, as a fraction of its radius:
// the mask is solid out to here and gives way in the sliver that is left, which
// is the edge's own anti-aliasing and nothing else.
//
// It was 0.75 — solid to three quarters, faded over the last quarter — and that
// fade is what the brush painted wrong. A spot aimed at something bigger than
// itself (the crown of a tree against flat sky) faded out ON the thing, so the
// thing came back at partial weight: a grey cloud the shape of the brush, and
// the fill showing through it where the mask had already given way. Measured
// (heal-repro2.mjs, one stroke of four spots over the crown, |after - truth|
// over the mask): 34.02 with the fade against 3.29 without; the outer band,
// 0.85r to r, 154.4 against 10.2; and the seam outside the brush 0 either way,
// because the fill is made of the rim's own light and already agrees with the
// photo along the edge — it needs no fade to hide behind.
export const HEAL_FEATHER = 0.995;
// out to 75% the mask is solid to completely eliminate defects, and the outer 25%
// is a smooth Hermite fade (smoothstep) to blend invisibly into the surrounding photo.
export const HEAL_FEATHER = 0.75;
// The finest of the fill's blurs, as a multiple of the widest brush in the
// stroke. The light of the rim has to reach the middle of the hole, so the
@@ -154,4 +143,22 @@ half4 main(float2 pos) {
// the same fraction of themselves.
export const HEAL_MASK_STOPS: number[] = [0, HEAL_FEATHER, 1];
// The same stops for the model's own pixels, and they are not the mask's. The
// model is the repair wherever it has one, so its shape is the whole spot and
// the last fraction of a radius is the only place it gives way. At
// HEAL_MASK_STOPS instead — the model solid out to 0.75 of the radius and faded
// out by the rim — the band between was the FILL's, and the fill is built from
// the frame's own light over the whole of its box (a stroke a third of the frame
// long is filled by blurs up to sixteen brush radii wide, so what it brings into
// that band is the box's average and not the rim's).
//
// Measured (probe-seam.mjs, one stroke of seven spots across the sky, `d` over
// the spots' own circles between the frame before the repair and after it): the
// band came back 8 to 11 levels below the photo it was made in — a ring, on the
// one line a repair is looked at along. With the model owning the band: 0.71
// levels of mean change over the whole repair, 8.5 at its worst pixel, and the
// edge strength at the rim after the repair 0.84 where the photo's own edges
// around it are 0.5 (1.52 before). Past 0.92 the numbers stop moving (0.73 /
// 8.7 / 0.88 at 0.995), so the model does own the band and what is left there is
// the shape of the spot, not the fill under it.
export const HEAL_PATCH_STOPS: number[] = [0, 0.92, 1];
docker/frontend/src/App.tsx
+8
View File
@@ -31,6 +31,7 @@ import { clearSession, forgetPhoto, loadPhoto, loadState, savePhoto, saveState,
import { analyzeCustomFrame, clearCustomFrame, loadCustomFrame, saveCustomFrame, type CustomFrameData } from './engine/customFrame';
import { keep } from './remember';
import { preloadSuperRes, upscaleJpeg } from './engine/superRes';
import { preloadInpaint, clearInpaintCache } from './engine/inpaint';
import { api, type Photo, type SavedRecipe } from './api';
import { refreshAuth, useAuth } from './auth';
import { TIERS } from './config/tiers';
@@ -1568,6 +1569,7 @@ export function Workspace() {
);
const clearHealSpots = useCallback(() => {
remember();
clearInpaintCache();
// The field is what "this photo was healed" means, so CLEAR drops it rather
// than leaving an empty list behind.
setAdjustment({ heal: undefined });
@@ -1616,6 +1618,7 @@ export function Workspace() {
const editBrushSpots = useCallback(
(spots: (HealSpot | MosaicSpot)[], undo: boolean) => {
if (undo) remember();
clearInpaintCache();
const list = spots.length ? spots : undefined;
if (brushTool === 'mosaic') setAdjustment({ mosaic: list as MosaicSpot[] | undefined });
else setAdjustment({ heal: list as HealSpot[] | undefined });
@@ -3707,6 +3710,11 @@ export function Workspace() {
setBrushTool(null);
} else {
setBrushTool('heal');
// The 27MB of model behind the brush, fetched while the visitor aims.
// Turning the tool on is not the repair, but it is the moment the
// repair becomes certain enough to spend the download on — see
// engine/inpaint.ts on why preloading is bytes and not the runtime.
preloadInpaint();
}
},
},
docker/frontend/src/engine/exportEngine.ts
+19 -20
View File
@@ -53,9 +53,10 @@ import {
HEAL_FILL_BLUR,
HEAL_FILL_MAX_LEVELS,
HEAL_MASK_STOPS,
HEAL_PATCH_STOPS,
readHeal,
} from '../../shared/utils/heal';
import { fastLocalInpaint, inpaintCropSize, INPAINT_REF_LONG } from './inpaint';
import { inpaintCrop, inpaintCropSize, INPAINT_REF_LONG } from './inpaint';
import { mosaicSkSL, mosaicUniforms, readMosaic } from '../../shared/utils/mosaic';
import { gradientMaskSkSL, maskUniforms, readMasks, masksHaveSpatial } from '../../shared/utils/gradientMask';
import {
@@ -1619,13 +1620,9 @@ function clusterHealSpots(
fillCanvas.drawImageRect(cutImage, box, onto, own(Skia.Paint()));
}
// The other half of the fill, and the better one: the pixels under the
// brush invented from the pixels around them (engine/inpaint.ts) rather
// than averaged from the rim — the rim's own light, diffused inward, with
// the rim's own grain over it. It is synchronous. It used to be MI-GAN
// over an ONNX runtime here, awaited; that was the lag and the worse
// repair both (heal.ts and inpaint.ts carry the numbers), so the model is
// gone and this is the rung it used to fall back on.
// The other half of the fill, and the better one where it can be had: the
// pixels under the brush invented by MI-GAN from the pixels around them
// (engine/inpaint.ts) rather than averaged from the rim.
//
// One drag is one repair, and a drag that runs the length of the frame is
// solved in crops of its own extent rather than in one downscaled box that
@@ -1654,12 +1651,13 @@ function clusterHealSpots(
const srcH = srcBottom - srcTop;
if (srcW >= 8 && srcH >= 8) {
// The crop's grid comes from its share of the FRAME and not from the
// pixels the frame is rendered at, so the same repair is the same
// repair at every zoom and at the export's own size — see
// The model's grid comes from the crop's share of the FRAME and not
// from the pixels the frame is rendered at, so the same repair is the
// same repair at every zoom and at the export's own size — see
// INPAINT_REF_LONG. It is also finer than it used to be wherever the
// stage is small: a crop a tenth of a 768px preview is 77 pixels, and
// 77 pixels of context invent a coarser repair than a bigger one.
// 77 pixels of context is what the model used to be asked to invent
// from.
const { w: mw, h: mh } = inpaintCropSize(
(srcW / width) * INPAINT_REF_LONG,
(srcH / height) * INPAINT_REF_LONG
@@ -1711,18 +1709,20 @@ function clusterHealSpots(
}
}
}
const patch = fastLocalInpaint(cropPixels, mw, mh, hole);
const patch = await inpaintCrop(cropPixels, mw, mh, hole).catch(() => null);
if (patch) {
const patchSurface = own(createSurface(mw, mh));
if (patchSurface && patchSurface.getCanvas().writePixels(patch, mw, mh, 0, 0)) {
flush(patchSurface);
const patchImage = own(patchSurface.makeImageSnapshot());
if (patchImage) {
// The invented pixels under the spot, and nowhere else at
// all. They are laid through the mask's own stops: they are
// the repair wherever they exist, so they own the whole spot
// and give way only in the sliver the mask's edge fades over,
// exactly as the blur under them does.
// The model's pixels under the spot, and nowhere else at
// all. What shape they are laid through is the join between
// the model's invention and the photo's real texture, and it
// is not the mask's shape: the model is the repair wherever
// it has one, so it owns the whole spot and gives way only in
// the sliver the mask itself fades over (HEAL_PATCH_STOPS,
// and the numbers that put it there are in heal.ts).
const patchSurface = own(createSurface(srcW, srcH));
if (patchSurface) {
const patchCanvas = patchSurface.getCanvas();
@@ -1745,7 +1745,7 @@ function clusterHealSpots(
[cx, cy],
rad,
new Float32Array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]),
HEAL_MASK_STOPS,
HEAL_PATCH_STOPS,
Skia.TileMode.Clamp
)
);
@@ -1781,7 +1781,6 @@ function clusterHealSpots(
}
}
}
flush(fillSurface);
const fillImage = own(fillSurface.makeImageSnapshot());
docker/frontend/src/engine/inpaint.ts
+392 -44
View File
@@ -1,50 +1,73 @@
// 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. What fills it now is a diffusion of that same rim,
// a harmonic field with the rim's own grain laid back over it (fastLocalInpaint
// below): still no pixel copied from anywhere, and the light of the photo rather
// than of the thing being repaired.
//
// It was MI-GAN over onnxruntime-web for a while too — 27MB of weights fetched
// and a session built on the visitor's machine, asked only for the pixels under
// the brush. It is gone: on the renderer's own probe (heal-repro2.mjs, five
// cases, the same mask and the same frames for both authors) the model was the
// slower half AND the worse one on every case it was given — crown-fit 9.28
// against 3.29 and 2273ms against 416ms, crown-wide 12.93/4.28 at 2279/359ms,
// sky-speck 0.72/0.34 at 4941/559ms, ridge-run 64.29/57.02 at 13013/1170ms —
// and on a flat rim it hands back the pixels it was given (migan-probe.mjs: 0.16
// of mean change over the plate it was asked to invent) after four seconds. A
// drag is solved on every move of the brush, so those seconds were the lag.
//
// ponytail: a diffusion, not a synthesis — over a hole wider than the structure
// around it there is nothing to continue and the answer is a plausible average.
// The rung for a synthesis is a model, and the numbers above are what one has to
// beat before it comes back.
//
// What the renderer hands over is a crop of the frame around the brush and the
// brush's own circles as a mask; what comes back is the same crop with the hole
// invented. Nothing outside the mask is read by the renderer (see step 6e), so a
// wrong corner of the crop costs nothing.
import type * as Ort from 'onnxruntime-web';
// The longest side a crop is ever answered at. A repair is a few brush widths
// across — the crop has all the context it needs long before this — and the
// renderer's own crop is cut to this grid, so the work is bounded by it rather
// than by the frame.
// 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 frame's long side, as a crop's grid is measured. The stage re-renders the
// frame on the zoom ladder (PREVIEW_QUANTUM in App.tsx, up to PREVIEW_HARD_MAX),
// so the same spots of the same photo are a different number of pixels at every
// zoom; a crop sized from those pixels would be a different grid each time, and
// every grid invents its own pixels — which is why a repair used to change every
// time the user zoomed back out of it, and drifted a little further with each
// pass. Sized against a fixed long side instead, the crop's box is a fraction of
// the frame and the grid it asks for is the same repair at every zoom and at the
// export's own size. 2048 is past the stage's own PREVIEW_MAX, so a crop small
// enough to matter is never scaled up to reach it.
// The frame's long side, as a crop's grid is measured for the model. The stage
// re-renders the frame on the zoom ladder (PREVIEW_QUANTUM in App.tsx, up to
// PREVIEW_HARD_MAX), so the same spots of the same photo are a different number
// of pixels at every zoom; a crop sized from those pixels handed the model a
// different grid each time, and every grid invents its own pixels — which is why
// a repair changed every time the user zoomed back out of it, and drifted a
// little further with each pass. Sized against a fixed long side instead, the
// crop's box is a fraction of the frame and the grid it asks for is the same
// repair at every zoom and at the export's own size. 2048 is past the stage's
// own PREVIEW_MAX, so a crop small enough to matter is never scaled up to reach
// it.
//
// ponytail: the grid is the same; the pixels handed to it are still read from
// whatever the stage rendered at, so the model's answer is the same repair to
// within the resampling of its input and not to the byte. Caching the patch
// across renders under a key that names the frame and the spots rather than the
// render's pixels is the rung that makes it identical, and it is the one to take
// if a zoom still shows a repair move.
export const INPAINT_REF_LONG = 2048;
// 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;
function isMobile(): boolean {
if (typeof navigator === 'undefined') return false;
@@ -55,6 +78,111 @@ function isMobile(): boolean {
);
}
function isIOS(): boolean {
if (typeof navigator === 'undefined') return false;
const ua = navigator.userAgent || '';
return (
/iPad|iPhone|iPod/i.test(ua) ||
(navigator.platform === 'MacIntel' && (navigator.maxTouchPoints || 0) > 1)
);
}
// 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 {
const in0 = new ort.Tensor('uint8', new Uint8Array(3 * WARM_SIDE * WARM_SIDE), [
1, 3, WARM_SIDE, WARM_SIDE,
]);
const in1 = new ort.Tensor('uint8', new Uint8Array(WARM_SIDE * WARM_SIDE).fill(255), [
1, 1, WARM_SIDE, WARM_SIDE,
]);
const out = await session.run({
[session.inputNames[0]]: in0,
[session.inputNames[1]]: in1,
});
try {
in0.dispose?.();
in1.dispose?.();
out[session.outputNames[0]]?.dispose?.();
} catch {}
return true;
} catch {
return false;
}
}
// Whether the browser has a device at all.
// Whether the browser has a device at all.
// Mobile devices (Android Chrome and iOS WebKit) have unified memory architectures where
// WebGPU buffer allocation destroys WebGL contexts (CanvasKit GrDirectContext loss -> black screen).
// Therefore WebGPU is strictly restricted to desktop.
async function hasDevice(): Promise<boolean> {
if (isMobile()) return false;
try {
return Boolean(await navigator.gpu?.requestAdapter());
} 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> {
if (await hasDevice()) {
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;
const mobile = isMobile();
// On mobile devices, 1 worker thread keeps memory strictly under the browser tab limit.
// On desktop, up to 4 threads if cross-origin isolated.
ort.env.wasm.numThreads = mobile
? 1
: self.crossOriginIsolated
? Math.min(4, navigator.hardwareConcurrency || 1)
: 1;
// On mobile and desktop browsers, proxy=false runs direct in-thread WASM execution,
// avoiding Web Worker proxy script loading failures and cross-origin worker restrictions.
ort.env.wasm.proxy = false;
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.
// Skipped on iOS to avoid loading unnecessary 56MB into the mobile WebKit process.
export function preloadInpaint(): void {
if (isIOS() || isMobile()) return;
void load().catch(() => {});
}
const byte = (v: number) => (v < 0 ? 0 : v > 255 ? 255 : Math.round(v));
// How many pixels the hole is.
@@ -325,6 +453,226 @@ export function fastLocalInpaint(
return out;
}
// What the model is handed: the crop at the model's own grid, and the mask in
// the model's own polarity — 255 where the frame is known, 0 over the hole.
//
// The pixels it reads UNDER the hole are the light of the hole's own rim, spread
// inward, and not the flat plate the crop used to be filled with: that was the
// mean of every pixel of the crop, which is a colour that appears nowhere near
// the repair as soon as the crop is bigger than the thing inside it (a blemish
// on a dark ground, a wire across a bright sky). The network then reads a hard
// step at the mask's own edge, answers in the plate's colour, and the repair
// lands as a patch of the wrong light on the frame's own. The rim's light meets
// the photo along the mask's edge by construction, so there is no step for the
// network to answer, and the context it reads is the photo's own.
//
// The mask is downsized by "any hole pixel in this cell" rather than by whichever
// pixel a cell's corner samples: the renderer hands a crop cut to this grid, and
// snapping that to whole strides of the graph can drop the last rows of it — with
// it, a brush thinner than the loss, and a frame handed to the model with nothing
// to invent comes back as the photo it already had.
export function cropPlate(
rgba: Uint8Array | Uint8ClampedArray,
width: number,
height: number,
hole: Uint8Array,
w: number,
h: number
): { src: Uint8Array; hole: Uint8Array } {
const sx = width / w;
const sy = height / h;
const small = new Uint8Array(4 * w * h);
for (let y = 0; y < h; y++) {
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) * 4;
small[d] = rgba[s];
small[d + 1] = rgba[s + 1];
small[d + 2] = rgba[s + 2];
small[d + 3] = 255;
}
}
const smallHole = new Uint8Array(w * h);
for (let i = 0; i < width * height; i++) {
if (!hole[i]) continue;
const x = Math.min(w - 1, Math.floor((i % width) / sx));
const y = Math.min(h - 1, Math.floor(((i / width) | 0) / sy));
smallHole[y * w + x] = 1;
}
const plane = w * h;
const src = new Uint8Array(3 * plane);
const mask = new Uint8Array(plane).fill(255);
const field = rimFill(small, w, h, smallHole, 0);
for (let d = 0; d < plane; d++) {
const s = d * 4;
if (smallHole[d]) {
src[d] = byte(field.r[d]);
src[plane + d] = byte(field.g[d]);
src[2 * plane + d] = byte(field.b[d]);
mask[d] = 0;
} else {
src[d] = small[s];
src[plane + d] = small[s + 1];
src[2 * plane + d] = small[s + 2];
}
}
return { src, hole: mask };
}
const patchCache = new Map<string, Uint8Array>();
let inFlightInpaint: Promise<Uint8Array | null> | null = null;
export function clearInpaintCache(): void {
patchCache.clear();
}
// 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> {
patchCache.clear();
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.
export async function inpaintCrop(
rgba: Uint8Array | Uint8ClampedArray,
width: number,
height: number,
hole: Uint8Array
): Promise<Uint8Array | null> {
const maxSide = isMobile() ? 256 : INPAINT_MAX_SIDE;
let w = Math.max(8, Math.round(width));
let h = Math.max(8, Math.round(height));
const scale = Math.min(1, maxSide / 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);
// Quick signature to avoid redundant re-inferences on subsequent renders
let holeCount = 0;
let holeHash = 0;
const step = Math.max(1, Math.floor(hole.length / 64));
for (let i = 0; i < hole.length; i += step) {
if (hole[i]) {
holeCount++;
holeHash = (holeHash * 31 + i) | 0;
}
}
let pixHash = 0;
const pixStep = Math.max(4, Math.floor(rgba.length / 64) * 4);
for (let i = 0; i < rgba.length; i += pixStep) {
pixHash = (pixHash * 33 + rgba[i] + rgba[i + 1]) | 0;
}
const key = `${w}x${h}_${holeCount}_${holeHash}_${pixHash}`;
if (patchCache.has(key)) {
return patchCache.get(key)!.slice();
}
// Single-flight lock: serialize inference executions to prevent WebAssembly memory spikes
if (inFlightInpaint) {
try {
await inFlightInpaint;
} catch {}
if (patchCache.has(key)) {
return patchCache.get(key)!.slice();
}
}
const runPromise = (async (): Promise<Uint8Array | null> => {
// NCHW, 0..255 RGB — what the model was trained to read — and the mask in its
// own polarity. Both are built by cropPlate, which is where the light the
// model reads under the hole comes from.
const { src, hole: mask } = cropPlate(rgba, width, height, hole, w, h);
try {
// Allow up to 15000ms for neural inpainting session to initialize and synthesize scene context
const sessionPromise = load();
const loadWithTimeout = Promise.race([
sessionPromise,
new Promise<never>((_, reject) => setTimeout(() => reject(new Error('Load timeout')), 15000)),
]);
const { ort, session } = await loadWithTimeout;
const in0 = new ort.Tensor('uint8', src, [1, 3, h, w]);
const in1 = new ort.Tensor('uint8', mask, [1, 1, h, w]);
const runWithTimeout = Promise.race([
session.run({
[session.inputNames[0]]: in0,
[session.inputNames[1]]: in1,
}),
new Promise<never>((_, reject) =>
setTimeout(() => reject(new Error('Inference timeout')), 15000)
),
]);
const out = await runWithTimeout;
const tensor = out[session.outputNames[0]];
const oplane = w * h;
const data = tensor.data as Uint8Array;
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;
}
try {
in0.dispose?.();
in1.dispose?.();
tensor.dispose?.();
for (const k in out) out[k]?.dispose?.();
} catch {}
if (patchCache.size > 50) {
const first = patchCache.keys().next().value;
if (first) patchCache.delete(first);
}
patchCache.set(key, patch.slice());
return patch;
} catch {
// Fallback: fast local inpaint with multi-step Poisson harmonic propagation.
// Do NOT cache the fallback in patchCache so that once neural inpainting finishes
// initializing, subsequent renders will execute neural context interpolation.
return fastLocalInpaint(rgba, width, height, hole);
}
})();
inFlightInpaint = runPromise;
try {
return await runPromise;
} finally {
if (inFlightInpaint === runPromise) {
inFlightInpaint = 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