diff --git a/docker/frontend/src/engine/inpaint.ts b/docker/frontend/src/engine/inpaint.ts index a10b5cd..7654c00 100644 --- a/docker/frontend/src/engine/inpaint.ts +++ b/docker/frontend/src/engine/inpaint.ts @@ -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 { 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 { - if (isMobile()) return false; try { return Boolean(await navigator.gpu?.requestAdapter()); } catch { @@ -129,39 +80,31 @@ async function hasDevice(): Promise { } } -// 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 { - 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 { 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 { 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 { 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 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 { + 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 { 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 => { - // 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 | 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((_, 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 | 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((_, 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))); diff --git a/docker/frontend/src/engine/inpaintWorker.ts b/docker/frontend/src/engine/inpaintWorker.ts new file mode 100644 index 0000000..07e97f6 --- /dev/null +++ b/docker/frontend/src/engine/inpaintWorker.ts @@ -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) => void) | null; + postMessage: (message: InpaintReply, transfer?: Transferable[]) => void; +}; + +scope.onmessage = async (event: MessageEvent) => { + const { id, src, mask, w, h } = event.data; + let in0: Ort.Tensor | null = null; + let in1: Ort.Tensor | null = null; + let out: Record | 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 {} + } +}; diff --git a/docker/frontend/vite.config.ts b/docker/frontend/vite.config.ts index 15d706d..aa74349 100644 --- a/docker/frontend/vite.config.ts +++ b/docker/frontend/vite.config.ts @@ -18,6 +18,9 @@ process.env.VITE_VERSION = version; // that for the CanvasKit adapter. Same trick the W0 spike used, now for the app. export default defineConfig({ plugins: [react()], + worker: { + format: 'es', + }, resolve: { // 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.