fix(heal): enable MI-GAN neural inpaint on mobile with bounded memory to prevent grey blob

This commit is contained in:
2026-10-16 09:39:15 +07:00
parent 455800c248
commit ce31e498d9
+27 -25
View File
@@ -133,6 +133,13 @@ async function hasDevice(): Promise<boolean> {
// 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> {
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'] });
@@ -142,7 +149,7 @@ async function openSession(ort: typeof Ort): Promise<Ort.InferenceSession> {
// Whatever went wrong, the wasm provider is the fallback and not a failure.
}
}
return ort.InferenceSession.create(MODEL_URL, { executionProviders: ['wasm'] });
return ort.InferenceSession.create(MODEL_URL, wasmOptions);
}
// Loads the runtime and the model once per page, and runs one warm-up pass so
@@ -160,9 +167,6 @@ function load(): Promise<Loaded> {
: self.crossOriginIsolated
? Math.min(4, navigator.hardwareConcurrency || 1)
: 1;
// On mobile and desktop browsers, proxy=false runs direct in-thread WASM execution,
// avoiding Web Worker proxy script loading failures and cross-origin worker restrictions.
ort.env.wasm.proxy = false;
ort.env.webgpu.powerPreference = 'high-performance';
const session = await openSession(ort);
return { ort, session };
@@ -177,9 +181,7 @@ function load(): Promise<Loaded> {
// Start fetching the model before the brush that needs it, so the visitor waits
// for the repair and not for the download. Called when HEAL is switched on.
// Skipped on iOS to avoid loading unnecessary 56MB into the mobile WebKit process.
export function preloadInpaint(): void {
if (isIOS() || isMobile()) return;
void load().catch(() => {});
}
@@ -554,15 +556,6 @@ export async function inpaintCrop(
height: number,
hole: Uint8Array
): Promise<Uint8Array | null> {
// On mobile browsers (iOS WebKit and Android), compiling the 27MB ONNX model and
// executing WASM tensors alongside CanvasKit exceeds tab memory limits (~384MB ceiling),
// causing iOS Jetsam process termination and browser page reload.
// On mobile, fast local inpaint provides instantaneous (<5ms) Poisson rim relaxation
// and grain synthesis using <1MB RAM, ensuring 100% reliable repair without tab crashes.
if (isIOS() || isMobile()) {
return fastLocalInpaint(rgba, width, height, hole);
}
const maxSide = isMobile() ? 256 : INPAINT_MAX_SIDE;
let w = Math.max(8, Math.round(width));
let h = Math.max(8, Math.round(height));
@@ -616,6 +609,9 @@ export async function inpaintCrop(
// 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();
@@ -624,8 +620,11 @@ export async function inpaintCrop(
new Promise<never>((_, reject) => setTimeout(() => reject(new Error('Load timeout')), 15000)),
]);
const { ort, session } = await loadWithTimeout;
const in0 = new ort.Tensor('uint8', src, [1, 3, h, w]);
const in1 = new ort.Tensor('uint8', mask, [1, 1, h, w]);
// 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({
@@ -637,7 +636,7 @@ export async function inpaintCrop(
),
]);
const out = await runWithTimeout;
out = await runWithTimeout;
const tensor = out[session.outputNames[0]];
const oplane = w * h;
const data = tensor.data as Uint8Array;
@@ -648,12 +647,6 @@ export async function inpaintCrop(
patch[i * 4 + 2] = data[2 * oplane + i];
patch[i * 4 + 3] = 255;
}
try {
in0.dispose?.();
in1.dispose?.();
tensor.dispose?.();
for (const k in out) out[k]?.dispose?.();
} catch {}
if (patchCache.size > 50) {
const first = patchCache.keys().next().value;
@@ -661,11 +654,20 @@ export async function inpaintCrop(
}
patchCache.set(key, patch.slice());
return patch;
} catch {
} 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?.();
}
} catch {}
}
})();