fix(heal): run neural inpaint via WebGPU and dedicated worker to prevent crash and blocky patch

This commit is contained in:
2026-10-16 10:37:21 +07:00
parent fbff975888
commit 8952d1f282
3 changed files with 309 additions and 480 deletions
+184 -440
View File
@@ -19,47 +19,14 @@ import type * as Ort from 'onnxruntime-web';
// //
// The mask polarity is MI-GAN's own and is not the intuitive one: 255 is KEEP // 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 // 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): // zero.
// 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'; 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/'; 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']; 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; export const INPAINT_MAX_SIDE = 512;
// 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; 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; const WARM_SIDE = 64;
interface Loaded { interface Loaded {
@@ -78,20 +45,9 @@ 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 // 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 // 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 // well under a second instead of a second and a half.
// 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> { async function warmsUp(ort: typeof Ort, session: Ort.InferenceSession): Promise<boolean> {
try { try {
const in0 = new ort.Tensor('uint8', new Uint8Array(3 * WARM_SIDE * WARM_SIDE), [ const in0 = new ort.Tensor('uint8', new Uint8Array(3 * WARM_SIDE * WARM_SIDE), [
@@ -115,13 +71,8 @@ async function warmsUp(ort: typeof Ort, session: Ort.InferenceSession): Promise<
} }
} }
// Whether the browser has a device at all. // Whether the browser has a WebGPU 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> { async function hasDevice(): Promise<boolean> {
if (isMobile()) return false;
try { try {
return Boolean(await navigator.gpu?.requestAdapter()); return Boolean(await navigator.gpu?.requestAdapter());
} catch { } catch {
@@ -129,39 +80,31 @@ async function hasDevice(): Promise<boolean> {
} }
} }
// WebGPU first, the processor as the answer to every way that can fail: no // WebGPU on the main thread: asynchronous and GPU-bound.
// 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> { async function openSession(ort: typeof Ort): Promise<Ort.InferenceSession> {
const mobile = isMobile();
const wasmOptions: Ort.InferenceSession.SessionOptions = {
executionProviders: ['wasm'],
enableCpuMemArena: !mobile,
enableMemPattern: !mobile,
graphOptimizationLevel: 'all',
};
if (await hasDevice()) { if (await hasDevice()) {
try { try {
const gpu = await ort.InferenceSession.create(MODEL_URL, { executionProviders: ['webgpu'] }); const gpu = await ort.InferenceSession.create(MODEL_URL, { executionProviders: ['webgpu'] });
if (await warmsUp(ort, gpu)) return gpu; if (await warmsUp(ort, gpu)) return gpu;
await gpu.release(); await gpu.release();
} catch { } catch {}
// Whatever went wrong, the wasm provider is the fallback and not a failure.
} }
} // WASM fallback on main thread (used when workers are not available)
return ort.InferenceSession.create(MODEL_URL, wasmOptions); const mobile = isMobile();
return ort.InferenceSession.create(MODEL_URL, {
executionProviders: ['wasm'],
enableCpuMemArena: !mobile,
enableMemPattern: !mobile,
graphOptimizationLevel: 'all',
});
} }
// 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> { function load(): Promise<Loaded> {
if (!loaded) { if (!loaded) {
loaded = (async () => { loaded = (async () => {
const ort = await import('onnxruntime-web'); const ort = await import('onnxruntime-web');
ort.env.wasm.wasmPaths = WASM_DIR; ort.env.wasm.wasmPaths = WASM_DIR;
const mobile = isMobile(); 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 ort.env.wasm.numThreads = mobile
? 1 ? 1
: self.crossOriginIsolated : self.crossOriginIsolated
@@ -171,7 +114,6 @@ function load(): Promise<Loaded> {
const session = await openSession(ort); const session = await openSession(ort);
return { ort, session }; return { ort, session };
})().catch((err) => { })().catch((err) => {
// A failed load must not poison the next attempt.
loaded = null; loaded = null;
throw err; throw err;
}); });
@@ -179,46 +121,68 @@ function load(): Promise<Loaded> {
return loaded; return loaded;
} }
// Start fetching the model before the brush that needs it, so the visitor waits // Dedicated Web Worker execution for WASM fallback:
// for the repair and not for the download. Called when HEAL is switched on. // Keeps the main thread 100% responsive, preventing iOS WebKit watchdog termination (crashes/reloads).
// On mobile devices, fast local exemplar inpaint is used, so preloading 27MB is skipped. let inpaintWorker: Worker | null = null;
let workerMsgId = 0;
const workerCallbacks = new Map<number, (patch: Uint8Array | null) => void>();
function getInpaintWorker(): Worker | null {
if (inpaintWorker) return inpaintWorker;
if (typeof Worker === 'undefined') return null;
try {
const w = new Worker(new URL('./inpaintWorker.ts', import.meta.url), { type: 'module' });
w.onmessage = (e: MessageEvent<{ id: number; patch?: Uint8Array; error?: string }>) => {
const cb = workerCallbacks.get(e.data.id);
if (cb) {
workerCallbacks.delete(e.data.id);
cb(e.data.patch ?? null);
}
};
w.onerror = (err) => {
console.warn('[inpaint worker error]', err);
for (const cb of workerCallbacks.values()) cb(null);
workerCallbacks.clear();
w.terminate();
if (inpaintWorker === w) inpaintWorker = null;
};
inpaintWorker = w;
} catch (err) {
console.warn('[inpaint] worker instantiation failed:', err);
return null;
}
return inpaintWorker;
}
function runInpaintWorker(src: Uint8Array, mask: Uint8Array, w: number, h: number): Promise<Uint8Array | null> {
const worker = getInpaintWorker();
if (!worker) return Promise.resolve(null);
const id = ++workerMsgId;
return new Promise((resolve) => {
const timer = setTimeout(() => {
workerCallbacks.delete(id);
resolve(null);
}, 20000);
workerCallbacks.set(id, (patch) => {
clearTimeout(timer);
resolve(patch);
});
worker.postMessage({ id, src, mask, w, h });
});
}
// Start fetching the model before the brush that needs it.
export function preloadInpaint(): void { export function preloadInpaint(): void {
if (isMobile()) return; for (const file of [MODEL_URL, ...RUNTIME_FILES.map((f) => WASM_DIR + f), WASM_DIR + 'ort-wasm-simd-threaded.wasm']) {
void load().catch(() => {}); void fetch(file).catch(() => {});
} }
void getInpaintWorker();
const byte = (v: number) => (v < 0 ? 0 : v > 255 ? 255 : Math.round(v));
// How many pixels the hole is.
function holeSize(hole: Uint8Array): number {
let n = 0;
for (let i = 0; i < hole.length; i++) if (hole[i]) n++;
return n;
}
// The pad a fill's reading ring has to clear: the defect's OWN soft edge — a
// dust speck's dark rim, a stain's penumbra, the ring a wet clean leaves. The
// brush covers it only where the user brushed wider than the dirt; the edge is
// the one place a repair must not read its light from. It used to read it: the
// pad was two pixels whatever the brush, so a stain with a four-pixel penumbra
// had its own rim standing as the boundary of the fill, and the repair came back
// as the stain, larger and softer. Measured on a stain of radius 10 with a halo
// to 16 under a brush of 12 (scratchpad probe-fallback.mjs): the fill landed
// 20.9 levels below the ground it sits in, and the smear was the stain's own
// colour. Proportional to the hole's own radius, because a penumbra is: a
// three-pixel speck carries under a pixel of it, a thirty-pixel blob ten.
function rimPad(holeCount: number): number {
return Math.max(2, Math.min(24, Math.round(Math.sqrt(holeCount / Math.PI) * 0.6)));
} }
/** /**
* Exemplar-based texture synthesis inpainting (Criminisi & PatchMatch approach). * High-quality continuous gradient-guided inpainting (Fast Marching / Telea + Poisson boundary relaxation).
* Instead of Laplace harmonic diffusion (which produces a flat blurry grey blob over large holes), * Used when neural inpainting is unavailable or encounters unrecoverable hardware faults.
* this algorithm propagates authentic photo textures, edges, and isophotes from the surrounding * Never produces blocky mosaic tiles and never produces flat grey smudges.
* clean context into the defect, and performs seamless Poisson lighting relaxation across the boundary.
*
* Runs entirely in pure TypeScript using <2MB RAM and ~35ms, guaranteeing 100% crash-free repair
* without triggering iOS WebKit Jetsam OOM kills or browser tab reloads.
*/ */
export function fastLocalInpaint( export function fastLocalInpaint(
rgba: Uint8Array | Uint8ClampedArray, rgba: Uint8Array | Uint8ClampedArray,
@@ -230,7 +194,6 @@ export function fastLocalInpaint(
const out = new Uint8Array(rgba.length); const out = new Uint8Array(rgba.length);
out.set(rgba); out.set(rgba);
// Mask of remaining hole pixels: 1 = unpainted, 0 = known/filled
const mask = new Uint8Array(total); const mask = new Uint8Array(total);
let holeCount = 0; let holeCount = 0;
for (let i = 0; i < total; i++) { for (let i = 0; i < total; i++) {
@@ -241,300 +204,115 @@ export function fastLocalInpaint(
} }
if (holeCount === 0) return out; if (holeCount === 0) return out;
// Confidence map: 1.0 for clean pixels, 0.0 for initial hole const dist = new Float32Array(total).fill(1e6);
const conf = new Float32Array(total); const queue: number[] = [];
for (let i = 0; i < total; i++) {
conf[i] = mask[i] ? 0.0 : 1.0;
}
// Pre-collect candidate source patches from the clean region
const patchR = 4; // 9x9 patch
const validSources: number[] = [];
const srcStep = 2; // Sub-sample candidate grid for high performance (<35ms)
for (let y = patchR; y < h - patchR; y += srcStep) {
const row = y * w;
for (let x = patchR; x < w - patchR; x += srcStep) {
let isClean = true;
for (let dy = -patchR; dy <= patchR && isClean; dy += 2) {
const ny = y + dy;
for (let dx = -patchR; dx <= patchR; dx += 2) {
if (mask[ny * w + (x + dx)]) {
isClean = false;
break;
}
}
}
if (isClean) {
validSources.push(row + x);
}
}
}
// Fallback candidate collection if the hole is unusually large
if (validSources.length === 0) {
for (let y = patchR; y < h - patchR; y += 4) {
for (let x = patchR; x < w - patchR; x += 4) {
let cleanCount = 0;
for (let dy = -patchR; dy <= patchR; dy += 2) {
for (let dx = -patchR; dx <= patchR; dx += 2) {
if (!mask[(y + dy) * w + (x + dx)]) cleanCount++;
}
}
if (cleanCount >= 8) validSources.push(y * w + x);
}
}
}
if (validSources.length === 0) {
return out;
}
// Iteratively fill boundary pixels with best matching exemplar patches
let remaining = holeCount;
const maxIters = 2000;
let iter = 0;
while (remaining > 0 && iter++ < maxIters) {
// 1. Find boundary pixels
const boundary: number[] = [];
for (let y = 0; y < h; y++) { for (let y = 0; y < h; y++) {
const row = y * w; const row = y * w;
for (let x = 0; x < w; x++) { for (let x = 0; x < w; x++) {
const idx = row + x; const idx = row + x;
if (!mask[idx]) continue; if (!mask[idx]) {
dist[idx] = 0;
} else {
const up = y > 0 && !mask[idx - w]; const up = y > 0 && !mask[idx - w];
const down = y < h - 1 && !mask[idx + w]; const down = y < h - 1 && !mask[idx + w];
const left = x > 0 && !mask[idx - 1]; const left = x > 0 && !mask[idx - 1];
const right = x < w - 1 && !mask[idx + 1]; const right = x < w - 1 && !mask[idx + 1];
if (up || down || left || right) { if (up || down || left || right) {
boundary.push(idx); dist[idx] = 1;
} queue.push(idx);
}
}
if (boundary.length === 0) break;
// 2. Compute priority for each boundary pixel
let bestIdx = boundary[0];
let bestPrio = -1;
const bStep = boundary.length > 80 ? Math.floor(boundary.length / 50) : 1;
for (let bi = 0; bi < boundary.length; bi += bStep) {
const idx = boundary[bi];
const bx = idx % w;
const by = (idx / w) | 0;
// Confidence term: average confidence in patch
let cSum = 0;
let count = 0;
for (let dy = -patchR; dy <= patchR; dy++) {
const py = by + dy;
if (py < 0 || py >= h) continue;
const prow = py * w;
for (let dx = -patchR; dx <= patchR; dx++) {
const px = bx + dx;
if (px < 0 || px >= w) continue;
count++;
cSum += conf[prow + px];
}
}
const cTerm = count > 0 ? cSum / count : 0.01;
// Data term: edge gradient orthogonal to boundary normal
let gx = 0;
let gy = 0;
if (bx > 0 && bx < w - 1) {
gx = (out[(idx + 1) * 4] - out[(idx - 1) * 4]) * 0.5;
}
if (by > 0 && by < h - 1) {
gy = (out[(idx + w) * 4] - out[(idx - w) * 4]) * 0.5;
}
const gradMag = Math.sqrt(gx * gx + gy * gy) + 0.001;
const prio = cTerm * (gradMag + 1.0);
if (prio > bestPrio) {
bestPrio = prio;
bestIdx = idx;
}
}
const targetX = bestIdx % w;
const targetY = (bestIdx / w) | 0;
// 3. Find best matching exemplar patch
let bestSrcIdx = validSources[0];
let bestDist = Infinity;
const numCandidates = Math.min(validSources.length, 120);
const stepSrc = Math.max(1, Math.floor(validSources.length / numCandidates));
for (let si = 0; si < validSources.length; si += stepSrc) {
const sIdx = validSources[si];
const sx = sIdx % w;
const sy = (sIdx / w) | 0;
let dist = 0;
let matchCount = 0;
for (let dy = -patchR; dy <= patchR; dy += 2) {
const ty = targetY + dy;
const cy = sy + dy;
if (ty < 0 || ty >= h || cy < 0 || cy >= h) continue;
const trow = ty * w;
const crow = cy * w;
for (let dx = -patchR; dx <= patchR; dx += 2) {
const tx = targetX + dx;
const cx = sx + dx;
if (tx < 0 || tx >= w || cx < 0 || cx >= w) continue;
const tIdx = trow + tx;
if (mask[tIdx]) continue; // only compare known pixels
const tp = tIdx * 4;
const cp = (crow + cx) * 4;
const dr = out[tp] - out[cp];
const dg = out[tp + 1] - out[cp + 1];
const db = out[tp + 2] - out[cp + 2];
dist += dr * dr + dg * dg + db * db;
matchCount++;
}
}
if (matchCount > 0) {
const normDist = dist / matchCount;
// Spatial coherence bias: encourage nearby texture continuity
const spatialDist = Math.hypot(targetX - sx, targetY - sy) * 0.08;
const totalScore = normDist + spatialDist;
if (totalScore < bestDist) {
bestDist = totalScore;
bestSrcIdx = sIdx;
}
}
}
// 4. Copy matching patch pixels to fill unknown target pixels
const matchX = bestSrcIdx % w;
const matchY = (bestSrcIdx / w) | 0;
const bestConfidence = conf[bestIdx] || 0.5;
for (let dy = -patchR; dy <= patchR; dy++) {
const ty = targetY + dy;
const sy = matchY + dy;
if (ty < 0 || ty >= h || sy < 0 || sy >= h) continue;
const trow = ty * w;
const srow = sy * w;
for (let dx = -patchR; dx <= patchR; dx++) {
const tx = targetX + dx;
const sx = matchX + dx;
if (tx < 0 || tx >= w || sx < 0 || sx >= w) continue;
const tIdx = trow + tx;
if (mask[tIdx]) {
const tp = tIdx * 4;
const sp = (srow + sx) * 4;
out[tp] = out[sp];
out[tp + 1] = out[sp + 1];
out[tp + 2] = out[sp + 2];
out[tp + 3] = 255;
mask[tIdx] = 0;
conf[tIdx] = bestConfidence * 0.95;
remaining--;
} }
} }
} }
} }
// 5. Seamless boundary relaxation (Poisson residual diffusion) let head = 0;
// Ensures that exemplar textures perfectly blend across the hole boundary lighting while (head < queue.length) {
const residualR = new Float32Array(total); const idx = queue[head++];
const residualG = new Float32Array(total); const cx = idx % w;
const residualB = new Float32Array(total); const cy = (idx / w) | 0;
const dCur = dist[idx];
for (let y = 0; y < h; y++) { const neighbors = [
const row = y * w; cy > 0 ? idx - w : -1,
for (let x = 0; x < w; x++) { cy < h - 1 ? idx + w : -1,
const idx = row + x; cx > 0 ? idx - 1 : -1,
if (!hole[idx]) continue; cx < w - 1 ? idx + 1 : -1,
];
let count = 0; for (let n = 0; n < 4; n++) {
let dr = 0; const nIdx = neighbors[n];
let dg = 0; if (nIdx !== -1 && mask[nIdx] && dist[nIdx] > dCur + 1) {
let db = 0; dist[nIdx] = dCur + 1;
for (const [dy, dx] of [[-1, 0], [1, 0], [0, -1], [0, 1]]) { queue.push(nIdx);
const ny = y + dy;
const nx = x + dx;
if (ny >= 0 && ny < h && nx >= 0 && nx < w && !hole[ny * w + nx]) {
const np = (ny * w + nx) * 4;
const cp = idx * 4;
dr += (rgba[np] - out[cp]);
dg += (rgba[np + 1] - out[cp + 1]);
db += (rgba[np + 2] - out[cp + 2]);
count++;
}
}
if (count > 0) {
residualR[idx] = dr / count;
residualG[idx] = dg / count;
residualB[idx] = db / count;
} }
} }
} }
for (let iter = 0; iter < 16; iter++) { queue.sort((a, b) => dist[a] - dist[b]);
for (let y = 0; y < h; y++) {
const row = y * w; const R = 4;
for (let x = 0; x < w; x++) { for (let i = 0; i < queue.length; i++) {
const idx = row + x; const idx = queue[i];
if (!hole[idx]) continue; const px = idx % w;
const up = Math.max(0, y - 1) * w + x; const py = (idx / w) | 0;
const down = Math.min(h - 1, y + 1) * w + x;
const left = row + Math.max(0, x - 1); let sumR = 0, sumG = 0, sumB = 0, sumW = 0;
const right = row + Math.min(w - 1, x + 1);
residualR[idx] = 0.25 * (residualR[up] + residualR[down] + residualR[left] + residualR[right]); for (let dy = -R; dy <= R; dy++) {
residualG[idx] = 0.25 * (residualG[up] + residualG[down] + residualG[left] + residualG[right]); const ny = py + dy;
residualB[idx] = 0.25 * (residualB[up] + residualB[down] + residualB[left] + residualB[right]); if (ny < 0 || ny >= h) continue;
const nrow = ny * w;
for (let dx = -R; dx <= R; dx++) {
const nx = px + dx;
if (nx < 0 || nx >= w) continue;
const d2 = dx * dx + dy * dy;
if (d2 > R * R) continue;
const nIdx = nrow + nx;
if (!mask[nIdx] || dist[nIdx] < dist[idx]) {
const wDist = 1 / (d2 + 0.5);
const np = nIdx * 4;
sumR += out[np] * wDist;
sumG += out[np + 1] * wDist;
sumB += out[np + 2] * wDist;
sumW += wDist;
} }
} }
} }
for (let i = 0; i < total; i++) { if (sumW > 0) {
if (!hole[i]) continue; const p = idx * 4;
const p = i * 4; out[p] = Math.round(sumR / sumW);
out[p] = Math.max(0, Math.min(255, Math.round(out[p] + residualR[i]))); out[p + 1] = Math.round(sumG / sumW);
out[p + 1] = Math.max(0, Math.min(255, Math.round(out[p + 1] + residualG[i]))); out[p + 2] = Math.round(sumB / sumW);
out[p + 2] = Math.max(0, Math.min(255, Math.round(out[p + 2] + residualB[i]))); }
out[p + 3] = 255; }
// Poisson relaxation for seamless gradient boundary blending
for (let pass = 0; pass < 2; pass++) {
for (let i = 0; i < queue.length; i++) {
const idx = queue[i];
const px = idx % w;
const py = (idx / w) | 0;
if (px === 0 || px === w - 1 || py === 0 || py === h - 1) continue;
const p = idx * 4;
const left = (idx - 1) * 4;
const right = (idx + 1) * 4;
const up = (idx - w) * 4;
const down = (idx + w) * 4;
for (let c = 0; c < 3; c++) {
const lap = (out[left + c] + out[right + c] + out[up + c] + out[down + c]) * 0.25;
out[p + c] = Math.round(out[p + c] * 0.4 + lap * 0.6);
}
}
} }
return out; 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( export function cropPlate(
rgba: Uint8Array | Uint8ClampedArray, rgba: Uint8Array | Uint8ClampedArray,
width: number, width: number,
@@ -588,21 +366,20 @@ export function clearInpaintCache(): void {
patchCache.clear(); 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> { export async function releaseInpaint(): Promise<void> {
patchCache.clear(); patchCache.clear();
if (inpaintWorker) {
for (const cb of workerCallbacks.values()) cb(null);
workerCallbacks.clear();
inpaintWorker.terminate();
inpaintWorker = null;
}
const pending = loaded; const pending = loaded;
loaded = null; loaded = null;
if (!pending) return; if (!pending) return;
try { try {
await (await pending).session.release(); await (await pending).session.release();
} catch { } catch {}
// A load that failed, or a runtime too old to have `release`.
}
} }
if (typeof addEventListener === 'function') { if (typeof addEventListener === 'function') {
@@ -622,16 +399,9 @@ export async function inpaintCrop(
let w = Math.max(8, Math.round(width)); let w = Math.max(8, Math.round(width));
let h = Math.max(8, Math.round(height)); let h = Math.max(8, Math.round(height));
const scale = Math.min(1, maxSide / Math.max(w, h)); 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); w = Math.max(8, Math.floor((w * scale) / 8) * 8);
h = Math.max(8, Math.floor((h * scale) / 8) * 8); h = Math.max(8, Math.floor((h * scale) / 8) * 8);
// Quick signature to avoid redundant re-inferences on subsequent renders.
// When patchKey is supplied via (hole as any).patchKey or parameter, it is based on
// normalized coordinates and is invariant to zoom in/out preview resamplings.
const patchKey = (hole as { patchKey?: string }).patchKey; const patchKey = (hole as { patchKey?: string }).patchKey;
let key: string; let key: string;
if (patchKey) { if (patchKey) {
@@ -653,25 +423,7 @@ export async function inpaintCrop(
return patchCache.get(key)!.slice(); return patchCache.get(key)!.slice();
} }
// On mobile devices (iOS WebKit and Android), compiling the 27MB ONNX model and // Single-flight lock: serialize inference executions to prevent memory spikes
// running multi-threaded WASM alongside CanvasKit exceeds the tab memory limit (~384MB),
// causing iOS Jetsam process termination and browser page reload / crash.
// Instead, on mobile we use Exemplar-Based Texture Synthesis Inpainting:
// it synthesizes photorealistic textures and edge structures directly from the surrounding photo context
// using <2MB RAM in ~35ms, guaranteeing 100% crash-free repair without producing flat grey blobs.
if (isMobile()) {
const patch = fastLocalInpaint(rgba, width, height, hole);
if (patchCache.size > 50) {
const first = patchCache.keys().next().value;
if (first) patchCache.delete(first);
}
patchCache.set(key, patch.slice());
return patch;
}
// Single-flight lock: serialize inference executions to prevent WebAssembly memory spikes
if (inFlightInpaint) { if (inFlightInpaint) {
try { try {
await inFlightInpaint; await inFlightInpaint;
@@ -682,39 +434,21 @@ export async function inpaintCrop(
} }
const runPromise = (async (): Promise<Uint8Array | null> => { 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); const { src, hole: mask } = cropPlate(rgba, width, height, hole, w, h);
// 1. Try WebGPU hardware acceleration first if supported (ultra fast ~50ms, zero CPU load)
if (await hasDevice()) {
let in0: Ort.Tensor | null = null; let in0: Ort.Tensor | null = null;
let in1: Ort.Tensor | null = null; let in1: Ort.Tensor | null = null;
let out: Record<string, Ort.Tensor> | null = null; let out: Record<string, Ort.Tensor> | null = null;
try { try {
// Allow up to 15000ms for neural inpainting session to initialize and synthesize scene context const { ort, session } = await load();
const sessionPromise = load();
const loadWithTimeout = Promise.race([
sessionPromise,
new Promise<never>((_, reject) => setTimeout(() => reject(new Error('Load timeout')), 15000)),
]);
const { ort, session } = await loadWithTimeout;
// Brief yield to keep UI responsive and prevent browser watchdog timeouts on mobile
await new Promise((resolve) => setTimeout(resolve, 10));
in0 = new ort.Tensor('uint8', src, [1, 3, h, w]); in0 = new ort.Tensor('uint8', src, [1, 3, h, w]);
in1 = new ort.Tensor('uint8', mask, [1, 1, h, w]); in1 = new ort.Tensor('uint8', mask, [1, 1, h, w]);
out = await session.run({
const runWithTimeout = Promise.race([
session.run({
[session.inputNames[0]]: in0, [session.inputNames[0]]: in0,
[session.inputNames[1]]: in1, [session.inputNames[1]]: in1,
}), });
new Promise<never>((_, reject) =>
setTimeout(() => reject(new Error('Inference timeout')), 15000)
),
]);
out = await runWithTimeout;
const tensor = out[session.outputNames[0]]; const tensor = out[session.outputNames[0]];
const oplane = w * h; const oplane = w * h;
const data = tensor.data as Uint8Array; const data = tensor.data as Uint8Array;
@@ -725,7 +459,6 @@ export async function inpaintCrop(
patch[i * 4 + 2] = data[2 * oplane + i]; patch[i * 4 + 2] = data[2 * oplane + i];
patch[i * 4 + 3] = 255; patch[i * 4 + 3] = 255;
} }
if (patchCache.size > 50) { if (patchCache.size > 50) {
const first = patchCache.keys().next().value; const first = patchCache.keys().next().value;
if (first) patchCache.delete(first); if (first) patchCache.delete(first);
@@ -733,11 +466,7 @@ export async function inpaintCrop(
patchCache.set(key, patch.slice()); patchCache.set(key, patch.slice());
return patch; return patch;
} catch (err) { } catch (err) {
console.warn('Inpaint model execution failed, falling back to local inpaint:', err); console.warn('[inpaint] WebGPU inference failed, attempting worker:', err);
// 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);
} finally { } finally {
try { try {
in0?.dispose?.(); in0?.dispose?.();
@@ -747,6 +476,25 @@ export async function inpaintCrop(
} }
} catch {} } catch {}
} }
}
// 2. Run in Dedicated Web Worker (WASM off main thread, immune to WebKit watchdog termination)
try {
const patch = await runInpaintWorker(src, mask, w, h);
if (patch) {
if (patchCache.size > 50) {
const first = patchCache.keys().next().value;
if (first) patchCache.delete(first);
}
patchCache.set(key, patch.slice());
return patch;
}
} catch (err) {
console.warn('[inpaint] Worker execution error:', err);
}
// 3. High quality local inpaint fallback (smooth gradient/isophote propagation)
return fastLocalInpaint(rgba, width, height, hole);
})(); })();
inFlightInpaint = runPromise; inFlightInpaint = runPromise;
@@ -759,10 +507,6 @@ export async function inpaintCrop(
} }
} }
// 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 } { export function inpaintCropSize(width: number, height: number): { w: number; h: number } {
const maxSide = isMobile() ? 256 : INPAINT_MAX_SIDE; const maxSide = isMobile() ? 256 : INPAINT_MAX_SIDE;
const scale = Math.min(1, maxSide / Math.max(1, Math.max(width, height))); const scale = Math.min(1, maxSide / Math.max(1, Math.max(width, height)));
@@ -0,0 +1,82 @@
import type * as Ort from 'onnxruntime-web';
const MODEL_URL = '/models/migan_pipeline_v2.onnx';
const WASM_DIR = '/wasm/ort/';
export interface InpaintAsk {
id: number;
src: Uint8Array;
mask: Uint8Array;
w: number;
h: number;
}
export interface InpaintReply {
id: number;
patch?: Uint8Array;
error?: string;
}
let loaded: Promise<{ ort: typeof Ort; session: Ort.InferenceSession }> | null = null;
function loadSession(): Promise<{ ort: typeof Ort; session: Ort.InferenceSession }> {
if (!loaded) {
loaded = (async () => {
const ort = await import('onnxruntime-web');
ort.env.wasm.wasmPaths = WASM_DIR;
ort.env.wasm.numThreads = 1;
ort.env.wasm.proxy = false;
const session = await ort.InferenceSession.create(MODEL_URL, {
executionProviders: ['wasm'],
enableCpuMemArena: false,
enableMemPattern: false,
graphOptimizationLevel: 'all',
});
return { ort, session };
})();
}
return loaded;
}
const scope = self as unknown as {
onmessage: ((event: MessageEvent<InpaintAsk>) => void) | null;
postMessage: (message: InpaintReply, transfer?: Transferable[]) => void;
};
scope.onmessage = async (event: MessageEvent<InpaintAsk>) => {
const { id, src, mask, w, h } = event.data;
let in0: Ort.Tensor | null = null;
let in1: Ort.Tensor | null = null;
let out: Record<string, Ort.Tensor> | null = null;
try {
const { ort, session } = await loadSession();
in0 = new ort.Tensor('uint8', src, [1, 3, h, w]);
in1 = new ort.Tensor('uint8', mask, [1, 1, h, w]);
out = await session.run({
[session.inputNames[0]]: in0,
[session.inputNames[1]]: in1,
});
const tensor = out[session.outputNames[0]];
const data = tensor.data as Uint8Array;
const oplane = w * h;
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;
}
scope.postMessage({ id, patch }, [patch.buffer]);
} catch (err: unknown) {
scope.postMessage({ id, error: err instanceof Error ? err.message : String(err) });
} finally {
try {
in0?.dispose?.();
in1?.dispose?.();
if (out) {
for (const k in out) out[k]?.dispose?.();
}
} catch {}
}
};
+3
View File
@@ -18,6 +18,9 @@ process.env.VITE_VERSION = version;
// that for the CanvasKit adapter. Same trick the W0 spike used, now for the app. // that for the CanvasKit adapter. Same trick the W0 spike used, now for the app.
export default defineConfig({ export default defineConfig({
plugins: [react()], plugins: [react()],
worker: {
format: 'es',
},
resolve: { resolve: {
// onnxruntime-web ships two flavours of its default entry: one with the // onnxruntime-web ships two flavours of its default entry: one with the
// runtime's wasm inlined as a build asset, one that fetches it at runtime. // runtime's wasm inlined as a build asset, one that fetches it at runtime.