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

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

Browse Source
This commit is contained in:
3dtours
2026-10-16 10:37:21 +07:00
parent fbff975888
commit 8952d1f282
3 changed files with 309 additions and 480 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/frontend/src/engine/inpaint.ts
+224 -480
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
// 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.
// zero.
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 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 {
@@ -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
// 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.
// well under a second instead of a second and a half.
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), [
@@ -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 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.
// Whether the browser has a WebGPU device at all.
async function hasDevice(): Promise<boolean> {
if (isMobile()) return false;
try {
return Boolean(await navigator.gpu?.requestAdapter());
} catch {
@@ -129,39 +80,31 @@ async function hasDevice(): Promise<boolean> {
}
}
// 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.
// WebGPU on the main thread: asynchronous and GPU-bound.
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()) {
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.
}
} catch {}
}
return ort.InferenceSession.create(MODEL_URL, wasmOptions);
// WASM fallback on main thread (used when workers are not available)
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> {
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
@@ -171,7 +114,6 @@ function load(): Promise<Loaded> {
const session = await openSession(ort);
return { ort, session };
})().catch((err) => {
// A failed load must not poison the next attempt.
loaded = null;
throw err;
});
@@ -179,46 +121,68 @@ function load(): Promise<Loaded> {
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.
// On mobile devices, fast local exemplar inpaint is used, so preloading 27MB is skipped.
// Dedicated Web Worker execution for WASM fallback:
// Keeps the main thread 100% responsive, preventing iOS WebKit watchdog termination (crashes/reloads).
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 {
if (isMobile()) return;
void load().catch(() => {});
}
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)));
for (const file of [MODEL_URL, ...RUNTIME_FILES.map((f) => WASM_DIR + f), WASM_DIR + 'ort-wasm-simd-threaded.wasm']) {
void fetch(file).catch(() => {});
}
void getInpaintWorker();
}
/**
* Exemplar-based texture synthesis inpainting (Criminisi & PatchMatch approach).
* Instead of Laplace harmonic diffusion (which produces a flat blurry grey blob over large holes),
* this algorithm propagates authentic photo textures, edges, and isophotes from the surrounding
* 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.
* High-quality continuous gradient-guided inpainting (Fast Marching / Telea + Poisson boundary relaxation).
* Used when neural inpainting is unavailable or encounters unrecoverable hardware faults.
* Never produces blocky mosaic tiles and never produces flat grey smudges.
*/
export function fastLocalInpaint(
rgba: Uint8Array | Uint8ClampedArray,
@@ -230,7 +194,6 @@ export function fastLocalInpaint(
const out = new Uint8Array(rgba.length);
out.set(rgba);
// Mask of remaining hole pixels: 1 = unpainted, 0 = known/filled
const mask = new Uint8Array(total);
let holeCount = 0;
for (let i = 0; i < total; i++) {
@@ -241,300 +204,115 @@ export function fastLocalInpaint(
}
if (holeCount === 0) return out;
// Confidence map: 1.0 for clean pixels, 0.0 for initial hole
const conf = new Float32Array(total);
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++) {
const row = y * w;
for (let x = 0; x < w; x++) {
const idx = row + x;
if (!mask[idx]) continue;
const up = y > 0 && !mask[idx - w];
const down = y < h - 1 && !mask[idx + w];
const left = x > 0 && !mask[idx - 1];
const right = x < w - 1 && !mask[idx + 1];
if (up || down || left || right) {
boundary.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)
// Ensures that exemplar textures perfectly blend across the hole boundary lighting
const residualR = new Float32Array(total);
const residualG = new Float32Array(total);
const residualB = new Float32Array(total);
const dist = new Float32Array(total).fill(1e6);
const queue: number[] = [];
for (let y = 0; y < h; y++) {
const row = y * w;
for (let x = 0; x < w; x++) {
const idx = row + x;
if (!hole[idx]) continue;
let count = 0;
let dr = 0;
let dg = 0;
let db = 0;
for (const [dy, dx] of [[-1, 0], [1, 0], [0, -1], [0, 1]]) {
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 (!mask[idx]) {
dist[idx] = 0;
} else {
const up = y > 0 && !mask[idx - w];
const down = y < h - 1 && !mask[idx + w];
const left = x > 0 && !mask[idx - 1];
const right = x < w - 1 && !mask[idx + 1];
if (up || down || left || right) {
dist[idx] = 1;
queue.push(idx);
}
}
if (count > 0) {
residualR[idx] = dr / count;
residualG[idx] = dg / count;
residualB[idx] = db / count;
}
}
let head = 0;
while (head < queue.length) {
const idx = queue[head++];
const cx = idx % w;
const cy = (idx / w) | 0;
const dCur = dist[idx];
const neighbors = [
cy > 0 ? idx - w : -1,
cy < h - 1 ? idx + w : -1,
cx > 0 ? idx - 1 : -1,
cx < w - 1 ? idx + 1 : -1,
];
for (let n = 0; n < 4; n++) {
const nIdx = neighbors[n];
if (nIdx !== -1 && mask[nIdx] && dist[nIdx] > dCur + 1) {
dist[nIdx] = dCur + 1;
queue.push(nIdx);
}
}
}
for (let iter = 0; iter < 16; iter++) {
for (let y = 0; y < h; y++) {
const row = y * w;
for (let x = 0; x < w; x++) {
const idx = row + x;
if (!hole[idx]) continue;
const up = Math.max(0, y - 1) * w + x;
const down = Math.min(h - 1, y + 1) * w + x;
const left = row + Math.max(0, x - 1);
const right = row + Math.min(w - 1, x + 1);
residualR[idx] = 0.25 * (residualR[up] + residualR[down] + residualR[left] + residualR[right]);
residualG[idx] = 0.25 * (residualG[up] + residualG[down] + residualG[left] + residualG[right]);
residualB[idx] = 0.25 * (residualB[up] + residualB[down] + residualB[left] + residualB[right]);
queue.sort((a, b) => dist[a] - dist[b]);
const R = 4;
for (let i = 0; i < queue.length; i++) {
const idx = queue[i];
const px = idx % w;
const py = (idx / w) | 0;
let sumR = 0, sumG = 0, sumB = 0, sumW = 0;
for (let dy = -R; dy <= R; dy++) {
const ny = py + dy;
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;
}
}
}
if (sumW > 0) {
const p = idx * 4;
out[p] = Math.round(sumR / sumW);
out[p + 1] = Math.round(sumG / sumW);
out[p + 2] = Math.round(sumB / sumW);
}
}
for (let i = 0; i < total; i++) {
if (!hole[i]) continue;
const p = i * 4;
out[p] = Math.max(0, Math.min(255, Math.round(out[p] + residualR[i])));
out[p + 1] = Math.max(0, Math.min(255, Math.round(out[p + 1] + residualG[i])));
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;
}
// 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,
@@ -588,21 +366,20 @@ 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();
if (inpaintWorker) {
for (const cb of workerCallbacks.values()) cb(null);
workerCallbacks.clear();
inpaintWorker.terminate();
inpaintWorker = null;
}
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`.
}
} catch {}
}
if (typeof addEventListener === 'function') {
@@ -622,16 +399,9 @@ export async function inpaintCrop(
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.
// 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;
let key: string;
if (patchKey) {
@@ -653,25 +423,7 @@ export async function inpaintCrop(
return patchCache.get(key)!.slice();
}
// On mobile devices (iOS WebKit and Android), compiling the 27MB ONNX model and
// 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
// Single-flight lock: serialize inference executions to prevent memory spikes
if (inFlightInpaint) {
try {
await inFlightInpaint;
@@ -682,71 +434,67 @@ export async function inpaintCrop(
}
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);
let in0: Ort.Tensor | null = null;
let in1: Ort.Tensor | null = null;
let out: Record<string, Ort.Tensor> | null = null;
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;
// 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]);
in1 = new ort.Tensor('uint8', mask, [1, 1, h, w]);
const runWithTimeout = Promise.race([
session.run({
// 1. Try WebGPU hardware acceleration first if supported (ultra fast ~50ms, zero CPU load)
if (await hasDevice()) {
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 load();
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,
}),
new Promise<never>((_, reject) =>
setTimeout(() => reject(new Error('Inference timeout')), 15000)
),
]);
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;
}
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 model execution failed, falling back to local inpaint:', 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 {
try {
in0?.dispose?.();
in1?.dispose?.();
if (out) {
for (const k in out) out[k]?.dispose?.();
});
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;
}
} catch {}
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] WebGPU inference failed, attempting worker:', err);
} finally {
try {
in0?.dispose?.();
in1?.dispose?.();
if (out) {
for (const k in out) out[k]?.dispose?.();
}
} 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;
@@ -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 } {
const maxSide = isMobile() ? 256 : INPAINT_MAX_SIDE;
const scale = Math.min(1, maxSide / Math.max(1, Math.max(width, height)));
docker/frontend/src/engine/inpaintWorker.ts
+82
View File
@@ -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 {}
}
};