fix(heal): enable MI-GAN neural inpaint on mobile with bounded memory to prevent grey blob
This commit is contained in:
@@ -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 {}
|
||||
}
|
||||
})();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user