diff --git a/docker/frontend/public/models/migan_pipeline_v2.onnx b/docker/frontend/public/models/migan_pipeline_v2.onnx new file mode 100644 index 0000000..3f75ade Binary files /dev/null and b/docker/frontend/public/models/migan_pipeline_v2.onnx differ diff --git a/docker/frontend/scripts/heal-check.mjs b/docker/frontend/scripts/heal-check.mjs index 5fcd3e6..70b2547 100644 --- a/docker/frontend/scripts/heal-check.mjs +++ b/docker/frontend/scripts/heal-check.mjs @@ -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 diff --git a/docker/frontend/shared/utils/heal.ts b/docker/frontend/shared/utils/heal.ts index 9599764..06e8407 100644 --- a/docker/frontend/shared/utils/heal.ts +++ b/docker/frontend/shared/utils/heal.ts @@ -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]; diff --git a/docker/frontend/src/App.tsx b/docker/frontend/src/App.tsx index c61101d..a7411fa 100644 --- a/docker/frontend/src/App.tsx +++ b/docker/frontend/src/App.tsx @@ -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(); } }, }, diff --git a/docker/frontend/src/engine/exportEngine.ts b/docker/frontend/src/engine/exportEngine.ts index d0a3ed6..fcd5d01 100644 --- a/docker/frontend/src/engine/exportEngine.ts +++ b/docker/frontend/src/engine/exportEngine.ts @@ -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()); diff --git a/docker/frontend/src/engine/inpaint.ts b/docker/frontend/src/engine/inpaint.ts index 7eae2b5..3bd2230 100644 --- a/docker/frontend/src/engine/inpaint.ts +++ b/docker/frontend/src/engine/inpaint.ts @@ -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 | 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 { + 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 { + 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 { + 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 { + 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(); +let inFlightInpaint: Promise | 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 { + 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 { + + 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 => { + // 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((_, 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((_, 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