RecipesCam/docker/frontend/src/engine/inpaint.ts

import type * as Ort from 'onnxruntime-web';

// FX > HEAL, the half that invents. The dust brush used to be filled with a
// blur of the photo's own rim — the light next to the hole, carried in, which is
// right for a speck of dust on a wall and thin over anything with structure: a
// hole across a shirt comes back as the average of a shirt.
//
// The fill is now a learned one. MI-GAN (andraniksargsyan/migan, the
// `migan_pipeline_v2.onnx` weights inpaint-web ships) is 27MB and runs through
// onnxruntime-web on the visitor's own machine, like the upscaler next door and
// for the same reason: the photo never leaves the device. It is only asked for
// what a GAN is actually good at — the pixels under the brush, and not one
// pixel more: everything outside the mask is the frame as rendered, put back by
// the composite in the renderer (see step 6e). A model that is wrong about a
// corner of the crop costs nothing, because that corner is never read.
//
// The runtime and the model are static files beside the app; nothing here
// changes what the server can do.
//
// 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.

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 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 {
  ort: typeof Ort;
  session: Ort.InferenceSession;
}

let loaded: Promise<Loaded> | null = null;

function isMobile(): boolean {
  if (typeof navigator === 'undefined') return false;
  const ua = navigator.userAgent || '';
  return (
    /iPad|iPhone|iPod|Android|Mobile/i.test(ua) ||
    (navigator.platform === 'MacIntel' && (navigator.maxTouchPoints || 0) > 1)
  );
}

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.
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), [
      1, 3, WARM_SIDE, WARM_SIDE,
    ]);
    const in1 = new ort.Tensor('uint8', new Uint8Array(WARM_SIDE * WARM_SIDE).fill(255), [
      1, 1, WARM_SIDE, WARM_SIDE,
    ]);
    const out = await session.run({
      [session.inputNames[0]]: in0,
      [session.inputNames[1]]: in1,
    });
    try {
      in0.dispose?.();
      in1.dispose?.();
      out[session.outputNames[0]]?.dispose?.();
    } catch {}
    return true;
  } catch {
    return false;
  }
}

// Whether the browser has a device at all.
// WebGPU on iOS (iPhone / iPad / WKWebView) is experimental and unstable;
// compiling WGSL or dispatching ONNX pipelines crashes WebKit's GPUProcess,
// causing iOS to silently reload the page. iOS is therefore excluded from WebGPU.
async function hasDevice(): Promise<boolean> {
  if (isIOS()) return false;
  try {
    return Boolean(await navigator.gpu?.requestAdapter());
  } catch {
    return false;
  }
}

// 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.
async function openSession(ort: typeof Ort): Promise<Ort.InferenceSession> {
  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.
    }
  }
  return ort.InferenceSession.create(MODEL_URL, { executionProviders: ['wasm'] });
}

// 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
        ? Math.min(4, navigator.hardwareConcurrency || 1)
        : 1;
      // ALWAYS proxy ONNX WASM execution into a Web Worker.
      // Running heavy WASM inferences on the main thread locks the UI loop ("not responding")
      // and triggers browser watchdog termination (crash / full reload).
      ort.env.wasm.proxy = true;
      ort.env.webgpu.powerPreference = 'high-performance';
      const session = await openSession(ort);
      return { ort, session };
    })().catch((err) => {
      // A failed load must not poison the next attempt.
      loaded = null;
      throw err;
    });
  }
  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.
//
// The bytes, and only the bytes: importing the runtime starts its thread pool —
// a Web Worker each, a copy of the 27MB binary each — and a pool that then
// belongs to the page until it dies, whether or not a stroke is ever painted.
// Switching a tool on is not the work; the download is, and the browser's cache
// is what keeps it.
export function preloadInpaint(): void {
  for (const file of [...RUNTIME_FILES.map((f) => WASM_DIR + f), MODEL_URL]) {
    void fetch(file).catch(() => {});
  }
}

interface CachedPatch {
  key: string;
  patch: Uint8Array;
}
let patchCache: CachedPatch | null = null;
let inFlightInpaint: Promise<Uint8Array | null> | null = null;

export function clearInpaintCache(): void {
  patchCache = null;
}

// 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 = 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`.
  }
}

if (typeof addEventListener === 'function') {
  addEventListener('pagehide', () => {
    void releaseInpaint();
  });
}

// The pixels under the brush, invented from the pixels around them.
//
// `rgba` is the crop the renderer cut out of the frame, at its own resolution;
// `hole` is one byte per pixel, non-zero for the pixels the brush painted. The
// answer is the same crop in RGBA at the same size, or null when the model could
// not be had — the caller keeps its blur fill in that case, and a repair that is
// merely soft beats no repair at all.
//
// A crop larger than INPAINT_MAX_SIDE on its longest side is scaled down to fit
// and the answer is scaled back up by the caller. A brush stroke is a soft shape
// made of soft shapes; what a factor of two of resample costs it is nothing
// beside the seconds it saves.
export async function inpaintCrop(
  rgba: Uint8Array | Uint8ClampedArray,
  width: number,
  height: number,
  hole: Uint8Array
): Promise<Uint8Array | null> {
  const maxSide = isMobile() ? 256 : INPAINT_MAX_SIDE;
  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
  let holeCount = 0;
  let holeHash = 0;
  const step = Math.max(1, Math.floor(hole.length / 64));
  for (let i = 0; i < hole.length; i += step) {
    if (hole[i]) {
      holeCount++;
      holeHash = (holeHash * 31 + i) | 0;
    }
  }
  let pixHash = 0;
  const pixStep = Math.max(4, Math.floor(rgba.length / 64) * 4);
  for (let i = 0; i < rgba.length; i += pixStep) {
    pixHash = (pixHash * 33 + rgba[i] + rgba[i + 1]) | 0;
  }
  const key = `${w}x${h}_${holeCount}_${holeHash}_${pixHash}`;

  if (patchCache && patchCache.key === key) {
    return patchCache.patch.slice();
  }

  // Single-flight lock: serialize inference executions to prevent WebAssembly memory spikes on mobile
  if (inFlightInpaint) {
    try {
      await inFlightInpaint;
    } catch {}
    if (patchCache && patchCache.key === key) {
      return patchCache.patch.slice();
    }
  }

  const runPromise = (async (): Promise<Uint8Array | null> => {
    // NCHW, 0..255 RGB — what the model was trained to read — and the mask in its
    // own polarity: 255 everywhere the frame is known, 0 over the hole.
    const src = new Uint8Array(3 * w * h);
    const mask = new Uint8Array(w * h).fill(255);
    const sx = width / w;
    const sy = height / h;
    const plane = w * h;
    for (let y = 0; y < h; y++) {
      // Nearest source pixel rather than a box filter: the crop is the frame at
      // its own resolution whenever it fits, so the two agree pixel for pixel and
      // this loop never runs at all in the common case.
      const syy = Math.min(height - 1, Math.floor(y * sy));
      for (let x = 0; x < w; x++) {
        const sxx = Math.min(width - 1, Math.floor(x * sx));
        const s = (syy * width + sxx) * 4;
        const d = y * w + x;
        src[d] = rgba[s];
        src[plane + d] = rgba[s + 1];
        src[2 * plane + d] = rgba[s + 2];
        if (hole[syy * width + sxx]) mask[d] = 0;
      }
    }

    try {
      const { ort, session } = await load();
      const in0 = new ort.Tensor('uint8', src, [1, 3, h, w]);
      const in1 = new ort.Tensor('uint8', mask, [1, 1, h, w]);
      const out = await session.run({
        [session.inputNames[0]]: in0,
        [session.inputNames[1]]: in1,
      });
      const tensor = out[session.outputNames[0]];
      const oplane = w * h;
      const data = tensor.data as Uint8Array;
      // Back to RGBA, opaque: the crop covers the frame, which has no transparency
      // of its own to carry.
      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;
      }
      try {
        in0.dispose?.();
        in1.dispose?.();
        tensor.dispose?.();
        for (const k in out) out[k]?.dispose?.();
      } catch {}

      patchCache = { key, patch: patch.slice() };
      return patch;
    } catch (err) {
      // A model that cannot run is a soft repair, not a broken one — and this one
      // line is what keeps the fallback diagnosable instead of invisible, because
      // the picture that comes out of it looks like a workable repair.
      console.warn('[inpaint] model unavailable, keeping the blurred fill', err);
      return null;
    }
  })();

  inFlightInpaint = runPromise;
  try {
    return await runPromise;
  } finally {
    if (inFlightInpaint === runPromise) {
      inFlightInpaint = null;
    }
  }
}

// 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)));
  return {
    w: Math.max(8, Math.floor((Math.max(1, Math.round(width)) * scale) / 8) * 8),
    h: Math.max(8, Math.floor((Math.max(1, Math.round(height)) * scale) / 8) * 8),
  };
}