web: EXPORT offers a size, and a bigger one is upscaled in the browser
The server still never sees a photo, so the model has to run in the page. Real-ESRGAN x4v3 ships as a 4.9MB ONNX in public/models and is loaded lazily on the first export that actually needs it; the wasm runtime is copied next to CanvasKit at build time and stays lazily fetched, cached for 30 days. Vite is told onnxruntime-web is external-wasm so no 28MB asset lands in the bundle. UNCHANGED keeps the old path and the tier cap; 2K/4K/custom upscale only when the request is larger than the photo being edited, otherwise they resize down. Guests keep UNCHANGED and 2K. Tiling is 256px with an 8px overlap, so memory follows the target size rather than four times it.
This commit is contained in:
@@ -0,0 +1,168 @@
|
||||
import type * as Ort from 'onnxruntime-web';
|
||||
|
||||
// Inventing pixels, in the browser. A small Real-ESRGAN — realesr-general-x4v3,
|
||||
// a 4.9MB SRVGGNetCompact that turns one pixel into four — runs through
|
||||
// onnxruntime-web on the visitor's own machine. The photo never leaves the
|
||||
// device, so the API's promise ("Never sees a photo") survives an export that
|
||||
// asks for more pixels than the picture has.
|
||||
//
|
||||
// Both the runtime (served next to canvaskit.wasm, see scripts/copy-wasm.mjs)
|
||||
// and the model (public/models) are static files; nothing here changes what the
|
||||
// server can do.
|
||||
|
||||
const MODEL_URL = '/models/realesr-general-x4v3.onnx';
|
||||
// The wasm dir, not the file: the runtime picks its own name inside it.
|
||||
const WASM_DIR = '/wasm/ort/';
|
||||
// The model's own factor is fixed at 4; the destination is whatever the export
|
||||
// asked for, so each finished tile is drawn to its place in the output rather
|
||||
// than assembled at 4x and shrunk (which would cost the memory of both).
|
||||
const MODEL_SCALE = 4;
|
||||
// Source pixels per pass, and the margin the model sees but the draw discards —
|
||||
// without it every tile boundary shows as a seam.
|
||||
const TILE = 256;
|
||||
const PAD = 8;
|
||||
|
||||
export interface UpscaleProgress {
|
||||
done: number;
|
||||
total: number;
|
||||
}
|
||||
|
||||
interface Loaded {
|
||||
ort: typeof Ort;
|
||||
session: Ort.InferenceSession;
|
||||
}
|
||||
|
||||
let loaded: Promise<Loaded> | null = null;
|
||||
|
||||
// Loads the runtime and the model once per page. WebGPU first — an export waits
|
||||
// seconds instead of a minute — and the same jsep build falls back to plain
|
||||
// WebAssembly when the browser has no GPU adapter to give.
|
||||
function load(): Promise<Loaded> {
|
||||
if (!loaded) {
|
||||
loaded = (async () => {
|
||||
const ort = await import('onnxruntime-web');
|
||||
ort.env.wasm.wasmPaths = WASM_DIR;
|
||||
// No cross-origin isolation on this site, so the threaded path has no
|
||||
// SharedArrayBuffer to use; saying so keeps the runtime from fetching a
|
||||
// proxy worker it cannot run.
|
||||
ort.env.wasm.numThreads = 1;
|
||||
const session = await ort.InferenceSession.create(MODEL_URL, { executionProviders: ['webgpu'] }).catch(() =>
|
||||
ort.InferenceSession.create(MODEL_URL, { executionProviders: ['wasm'] })
|
||||
);
|
||||
return { ort, session };
|
||||
})().catch((err) => {
|
||||
// A failed load must not poison the next attempt.
|
||||
loaded = null;
|
||||
throw err;
|
||||
});
|
||||
}
|
||||
return loaded;
|
||||
}
|
||||
|
||||
// The photo's bytes, enlarged so its longest edge is `targetLongest`. The caller
|
||||
// only gets here when the target is larger than the picture (App decides), and
|
||||
// gets back the same bytes untouched when it is not.
|
||||
export async function upscaleJpeg(
|
||||
bytes: Uint8Array,
|
||||
targetLongest: number,
|
||||
onProgress?: (p: UpscaleProgress) => void
|
||||
): Promise<Uint8Array> {
|
||||
const bitmap = await createImageBitmap(new Blob([bytes as BlobPart]));
|
||||
try {
|
||||
const w = bitmap.width;
|
||||
const h = bitmap.height;
|
||||
const longest = Math.max(w, h);
|
||||
if (targetLongest <= longest) return bytes;
|
||||
const scale = targetLongest / longest;
|
||||
|
||||
const srcCanvas = new OffscreenCanvas(w, h);
|
||||
const srcCtx = srcCanvas.getContext('2d', { willReadFrequently: true });
|
||||
const dstCanvas = new OffscreenCanvas(Math.max(1, Math.round(w * scale)), Math.max(1, Math.round(h * scale)));
|
||||
const dstCtx = dstCanvas.getContext('2d');
|
||||
if (!srcCtx || !dstCtx) return bytes;
|
||||
srcCtx.drawImage(bitmap, 0, 0);
|
||||
const src = srcCtx.getImageData(0, 0, w, h);
|
||||
|
||||
const { ort, session } = await load();
|
||||
const inputName = session.inputNames[0];
|
||||
const cols = Math.ceil(w / TILE);
|
||||
const rows = Math.ceil(h / TILE);
|
||||
let done = 0;
|
||||
onProgress?.({ done, total: cols * rows });
|
||||
|
||||
for (let ty = 0; ty < rows; ty++) {
|
||||
for (let tx = 0; tx < cols; tx++) {
|
||||
const x0 = tx * TILE;
|
||||
const y0 = ty * TILE;
|
||||
const tw = Math.min(TILE, w - x0);
|
||||
const th = Math.min(TILE, h - y0);
|
||||
// The margin the model gets: full on the inside, clipped at the photo's
|
||||
// own edge, so the tensor covers whole pixels only.
|
||||
const left = Math.min(PAD, x0);
|
||||
const top = Math.min(PAD, y0);
|
||||
const pw = tw + left + Math.min(PAD, w - (x0 + tw));
|
||||
const ph = th + top + Math.min(PAD, h - (y0 + th));
|
||||
|
||||
// NCHW, 0..1 RGB — what the model was trained to read.
|
||||
const input = new Float32Array(3 * pw * ph);
|
||||
const plane = pw * ph;
|
||||
for (let y = 0; y < ph; y++) {
|
||||
const srow = ((y0 - top + y) * w + (x0 - left)) * 4;
|
||||
for (let x = 0; x < pw; x++) {
|
||||
const s = srow + x * 4;
|
||||
input[y * pw + x] = src.data[s] / 255;
|
||||
input[plane + y * pw + x] = src.data[s + 1] / 255;
|
||||
input[2 * plane + y * pw + x] = src.data[s + 2] / 255;
|
||||
}
|
||||
}
|
||||
|
||||
const out = await session.run({ [inputName]: new ort.Tensor('float32', input, [1, 3, ph, pw]) });
|
||||
const tensor = out[session.outputNames[0]];
|
||||
const dims = tensor.dims as number[];
|
||||
const ow = dims[3];
|
||||
const oh = dims[2];
|
||||
const opts = tensor.data as Float32Array;
|
||||
|
||||
// Paint the tile's own pixels into a canvas of their own, then draw only
|
||||
// the unpadded middle of it into the output, at the output's scale.
|
||||
const tile = new OffscreenCanvas(ow, oh);
|
||||
const tileCtx = tile.getContext('2d');
|
||||
if (!tileCtx) return bytes;
|
||||
const img = tileCtx.createImageData(ow, oh);
|
||||
const oplane = ow * oh;
|
||||
for (let i = 0; i < oplane; i++) {
|
||||
const d = i * 4;
|
||||
img.data[d] = clamp8(opts[i]);
|
||||
img.data[d + 1] = clamp8(opts[oplane + i]);
|
||||
img.data[d + 2] = clamp8(opts[2 * oplane + i]);
|
||||
img.data[d + 3] = 255;
|
||||
}
|
||||
tileCtx.putImageData(img, 0, 0);
|
||||
dstCtx.drawImage(
|
||||
tile,
|
||||
left * MODEL_SCALE,
|
||||
top * MODEL_SCALE,
|
||||
tw * MODEL_SCALE,
|
||||
th * MODEL_SCALE,
|
||||
x0 * scale,
|
||||
y0 * scale,
|
||||
tw * scale,
|
||||
th * scale
|
||||
);
|
||||
|
||||
done++;
|
||||
onProgress?.({ done, total: cols * rows });
|
||||
}
|
||||
}
|
||||
|
||||
const blob = await dstCanvas.convertToBlob({ type: 'image/jpeg', quality: 0.95 });
|
||||
return new Uint8Array(await blob.arrayBuffer());
|
||||
} finally {
|
||||
bitmap.close();
|
||||
}
|
||||
}
|
||||
|
||||
function clamp8(v: number): number {
|
||||
const n = v * 255;
|
||||
return n <= 0 ? 0 : n >= 255 ? 255 : Math.round(n);
|
||||
}
|
||||
Reference in New Issue
Block a user