oidn-web 0.4.0 → 0.5.0

package/lib/UNet.d.ts

@@ -1,8 +1,8 @@
 import { HostTensor } from './tza';
-import { Tile } from './process';
+import { Tile, type HDRTransfer } from './process';
 import { type DynamicTileSetting } from './tileScheduler';
 import { type UNetModelSpec } from './modelSpec';
-import { type NativeUNetKernelSetting, type NativeUNetPrecisionSetting } from './nativeUNet';
+import { type NativeUNetGemmOptions, type NativeUNetKernelSetting, type NativeUNetPrecisionSetting } from './nativeUNet';
 export type UNetEngineSetting = 'auto' | 'wgsl' | 'webnn';
 interface HDRImageData {
     data: Float32Array;
@@ -22,9 +22,12 @@ interface GPUImageDataOutput {
 export interface UNetExecutionStats {
     width: number;
     height: number;
-    tileWidth: number;
-    tileHeight: number;
     tileCount: number;
+    tileColumns: number;
+    tileRows: number;
+    tileOverlap: number;
+    inputPixelCount: number;
+    inputShapeCount: number;
     durationMs: number;
     tileTimeMs: {
         min: number;
@@ -35,12 +38,9 @@ export interface UNetExecutionStats {
 }
 declare class UNet {
     private _device;
-    private _tileWidth;
-    private _tileHeight;
-    private _tileOverlapX;
-    private _tileOverlapY;
     private _aux;
     private _hdr;
+    private _hdrTransfer;
     private _dataProcessGPU?;
     private _nativeExecutor?;
     private _webNNExecutor?;
@@ -49,10 +49,11 @@ declare class UNet {
     private _engine;
     private _dynamicTileController;
     private _lastExecution?;
-    constructor(hostTensors: Map<string, HostTensor>, backend: {
-        device: GPUDevice;
-        adapterInfo: GPUAdapterInfo;
-    }, opts?: {
+    private _activeExecutionFailures;
+    private _deviceLostObserved;
+    private _deviceLostSettled;
+    private _deviceLostReason;
+    constructor(hostTensors: Map<string, HostTensor>, device: GPUDevice, opts?: {
         /**
          * If use auxiliary data.
          */
@@ -61,26 +62,38 @@ declare class UNet {
          * If input is HDR image.
          */
         hdr?: boolean;
+        /** HDR transfer function expected by the trained model. */
+        hdrTransfer?: HDRTransfer;
         maxTileSize?: number;
         dynamicTile?: DynamicTileSetting;
-        /** Reserved for explicit native WGSL selection. */
+        /** Native WGSL or the experimental WebNN backend. */
         engine?: UNetEngineSetting;
-        /** Arithmetic/storage precision used by the native WGSL engine. */
+        /** Arithmetic/storage precision used by the native WGSL executor. */
         precision?: NativeUNetPrecisionSetting;
         /** Model-independent convolution kernel selection. */
         kernel?: NativeUNetKernelSetting;
+        gemm?: NativeUNetGemmOptions;
         /** Explicit descriptor for a new OIDN topology not in the built-in registry. */
         modelSpec?: UNetModelSpec;
     });
-    getDevice(): GPUDevice | undefined;
+    getDevice(): GPUDevice;
     /** Completes backend compilation before first interactive use. */
     prepare(): Promise<void>;
+    /**
+     * Prepares the input shapes selected for an image before its first denoise.
+     * Hosts can call this while they still display their model-loading state.
+     */
+    prepareForImage(width: number, height: number, options?: {
+        tileOverlap?: number;
+        wholeImage?: boolean;
+    }): Promise<void>;
     getRuntimeInfo(): {
         configuredEngine: UNetEngineSetting;
         gpuEngine: "wgsl" | "webnn";
         precision: import("./nativeUNet").NativeUNetPrecision;
         kernel: {
             configured: NativeUNetKernelSetting;
+            gemm: Readonly<Required<NativeUNetGemmOptions>>;
             maxSpatialInputBlocks: number;
             subgroupsAvailable: boolean;
         } | undefined;
@@ -89,6 +102,7 @@ declare class UNet {
         model: string;
         modelFamily: (string & {}) | "oidn-unet-small" | "oidn-unet-large";
         inputChannels: number;
+        hdrTransfer: HDRTransfer;
         dynamicTile: {
             enabled: boolean;
             currentTileSize: number;
@@ -97,17 +111,18 @@ declare class UNet {
             targetTileTimeMs: number;
         };
         lastExecution: UNetExecutionStats | undefined;
+        activeExecutionCount: number;
     };
+    private _observeDeviceLoss;
+    private _registerExecutionFailure;
     /** Captures per-node GPU timestamps for the next native tile execution. */
     profileNextExecution(): boolean;
     getLastExecutionProfile(): Promise<import("./nativeUNet").NativeUNetExecutionProfile> | undefined;
-    private _updateModel;
-    private _getTileSizeWithOverlap;
     private _processImageData;
     private _readTile;
     private _writeTile;
     private _executeTile;
-    tileExecute<T extends ImageData | HDRImageData | GPUImageData>({ color, albedo, normal, done, progress, denoiseAlpha }: {
+    tileExecute<T extends ImageData | HDRImageData | GPUImageData>({ color, albedo, normal, done, progress, denoiseAlpha, tileOverlap, wholeImage, scheduling, error }: {
         color: T;
         albedo?: ImageData | GPUImageData;
         normal?: ImageData | GPUImageData;
@@ -115,8 +130,27 @@ declare class UNet {
          * If denoise alpha channel. Otherwise denoise RGB channels.
          */
         denoiseAlpha?: boolean;
-        done: (outputData: T extends GPUImageData ? GPUImageDataOutput : T) => void;
-        progress?: (outputData: T extends GPUImageData ? GPUImageDataOutput : T, tileData: (T extends GPUImageData ? GPUImageDataOutput : T) | undefined, tile: Tile, currentIdx: number, totalIdx: number) => void;
+        /**
+         * Execute the complete input image as one tile, ignoring `maxTileSize`.
+         * The image must fit the device's buffer and dispatch limits.
+         */
+        wholeImage?: boolean;
+        /**
+         * Per-side context for boundaries shared with another tile. Defaults to
+         * half of the model receptive field rounded up to 16 pixels.
+         */
+        tileOverlap?: number;
+        /**
+         * How JavaScript yields between completed GPU tiles. `event-loop`
+         * (default) continues on the next macrotask. `animation-frame` waits for
+         * the next display frame, bounded by a short timer so hidden pages still
+         * complete.
+         */
+        scheduling?: 'animation-frame' | 'event-loop';
+        done: (outputData: T extends GPUImageData ? GPUImageDataOutput : T) => void | Promise<void>;
+        /** Receives asynchronous execution, queue, and callback failures. */
+        error?: (reason: unknown) => void | Promise<void>;
+        progress?: (outputData: T extends GPUImageData ? GPUImageDataOutput : T, tileData: (T extends GPUImageData ? GPUImageDataOutput : T) | undefined, tile: Tile, currentIdx: number, totalIdx: number) => void | Promise<void>;
     }): () => void;
     dispose(): void;
 }

package/lib/UNet.js

@@ -1,8 +1,10 @@
 import { GPUDataProcess, Tile, avgLogLum, hdrTransferFuncCPU, hdrTransferFuncInverseCPU } from './process';
-import { DynamicTileController, fitTileDimension, OIDN_TILE_ALIGNMENT, waitForSubmittedGPUWork } from './tileScheduler';
+import { DynamicTileController, planTileGrid, OIDN_TILE_ALIGNMENT } from './tileScheduler';
 import { detectUNetModelSpec, validateUNetModel } from './modelSpec';
 import { NativeUNetExecutor } from './nativeUNet';
 import { WebNNUNetExecutor } from './webnnUNet';
+/** Upper bound on waiting for a display frame between tiles. */
+const ANIMATION_FRAME_FALLBACK_MS = 100;
 function roundUp(a, b) {
     return Math.ceil(a / b) * b;
 }
@@ -11,14 +13,9 @@ function isGPUImageData(data) {
 }
 class UNet {
     _device;
-    // TODO calculate the tile size from memory size
-    // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/core/unet_filter.cpp#L287
-    _tileWidth = 0;
-    _tileHeight = 0;
-    _tileOverlapX = 0;
-    _tileOverlapY = 0;
     _aux;
     _hdr;
+    _hdrTransfer;
     _dataProcessGPU;
     _nativeExecutor;
     _webNNExecutor;
@@ -27,9 +24,14 @@ class UNet {
     _engine;
     _dynamicTileController;
     _lastExecution;
-    constructor(hostTensors, backend, opts = {}) {
+    _activeExecutionFailures = new Set();
+    _deviceLostObserved = false;
+    _deviceLostSettled = false;
+    _deviceLostReason;
+    constructor(hostTensors, device, opts = {}) {
         this._aux = opts.aux || false;
         this._hdr = opts.hdr || false;
+        this._hdrTransfer = opts.hdrTransfer ?? 'pu';
         this._engine = opts.engine ?? 'auto';
         const modelSpec = opts.modelSpec ?? detectUNetModelSpec(hostTensors);
         const validatedModel = validateUNetModel(hostTensors, modelSpec);
@@ -41,12 +43,13 @@ class UNet {
                 `but aux=${this._aux} provides ${expectedInputChannels}`);
         }
         this._dynamicTileController = new DynamicTileController(opts.maxTileSize ?? 512, opts.dynamicTile);
-        this._device = backend.device;
+        this._device = device;
+        this._observeDeviceLoss();
         if (this._engine === 'webnn') {
             this._webNNExecutor = new WebNNUNetExecutor(this._device, validatedModel, { precision: opts.precision });
         }
         else {
-            this._nativeExecutor = new NativeUNetExecutor(this._device, validatedModel, { precision: opts.precision, kernel: opts.kernel });
+            this._nativeExecutor = new NativeUNetExecutor(this._device, validatedModel, { precision: opts.precision, kernel: opts.kernel, gemm: opts.gemm });
         }
     }
     getDevice() {
@@ -69,6 +72,25 @@ class UNet {
         }
         await this._nativeExecutor.prepare();
     }
+    /**
+     * Prepares the input shapes selected for an image before its first denoise.
+     * Hosts can call this while they still display their model-loading state.
+     */
+    async prepareForImage(width, height, options = {}) {
+        const defaultTileOverlap = roundUp(this._modelSpec.receptiveField / 2, OIDN_TILE_ALIGNMENT);
+        const resolvedTileOverlap = options.tileOverlap === undefined
+            ? defaultTileOverlap
+            : roundUp(Math.max(0, options.tileOverlap), OIDN_TILE_ALIGNMENT);
+        const plan = planTileGrid(width, height, options.wholeImage
+            ? roundUp(Math.max(width, height), OIDN_TILE_ALIGNMENT)
+            : this._dynamicTileController.tileSize, resolvedTileOverlap);
+        const shapes = [...new Map(plan.tiles.map(({ input }) => [
+                `${input.width}x${input.height}`,
+                { width: input.width, height: input.height }
+            ])).values()];
+        await this._webNNExecutor?.prewarm(shapes);
+        this._nativeExecutor?.prewarm(shapes);
+    }
     getRuntimeInfo() {
         return {
             configuredEngine: this._engine,
@@ -77,6 +99,7 @@ class UNet {
             kernel: this._nativeExecutor
                 ? {
                     configured: this._nativeExecutor.kernelSetting,
+                    gemm: this._nativeExecutor.gemm,
                     maxSpatialInputBlocks: this._nativeExecutor.maxSpatialInputBlocks,
                     subgroupsAvailable: this._nativeExecutor.subgroupsAvailable
                 }
@@ -86,6 +109,7 @@ class UNet {
             model: this._modelSpec.id,
             modelFamily: this._modelSpec.family,
             inputChannels: this._inputChannels,
+            hdrTransfer: this._hdrTransfer,
             dynamicTile: {
                 enabled: this._dynamicTileController.enabled,
                 currentTileSize: this._dynamicTileController.tileSize,
@@ -93,8 +117,39 @@ class UNet {
                 maxTileSize: this._dynamicTileController.maxTileSize,
                 targetTileTimeMs: this._dynamicTileController.targetTileTimeMs
             },
-            lastExecution: this._lastExecution
+            lastExecution: this._lastExecution,
+            activeExecutionCount: this._activeExecutionFailures.size
+        };
+    }
+    _observeDeviceLoss() {
+        // Object.create-based embedders/tests can bypass field initializers.
+        this._activeExecutionFailures ??= new Set();
+        if (this._deviceLostObserved)
+            return;
+        this._deviceLostObserved = true;
+        const deviceLost = this._device.lost;
+        if (!deviceLost)
+            return;
+        const failAll = (reason) => {
+            if (this._deviceLostSettled)
+                return;
+            this._deviceLostSettled = true;
+            this._deviceLostReason = reason;
+            const active = [...this._activeExecutionFailures];
+            this._activeExecutionFailures.clear();
+            for (const fail of active)
+                fail(reason);
         };
+        void deviceLost.then((info) => failAll(new Error(`WebGPU device lost: ${info.message}`)), failAll);
+    }
+    _registerExecutionFailure(fail) {
+        this._observeDeviceLoss();
+        this._activeExecutionFailures.add(fail);
+        if (this._deviceLostSettled) {
+            this._activeExecutionFailures.delete(fail);
+            queueMicrotask(() => fail(this._deviceLostReason));
+        }
+        return () => this._activeExecutionFailures.delete(fail);
     }
     /** Captures per-node GPU timestamps for the next native tile execution. */
     profileNextExecution() {
@@ -103,43 +158,6 @@ class UNet {
     getLastExecutionProfile() {
         return this._nativeExecutor?.getLastExecutionProfile();
     }
-    _updateModel(width, height) {
-        const maxTileSize = this._dynamicTileController.tileSize;
-        let tileWidth = fitTileDimension(width, maxTileSize);
-        let tileHeight = fitTileDimension(height, maxTileSize);
-        const defaultTileOverlap = roundUp(this._modelSpec.receptiveField / 2, OIDN_TILE_ALIGNMENT);
-        let tileOverlapX = defaultTileOverlap;
-        let tileOverlapY = defaultTileOverlap;
-        if (width <= maxTileSize) {
-            tileOverlapX = 0;
-        }
-        if (height <= maxTileSize) {
-            tileOverlapY = 0;
-        }
-        // Force width and height has same size. reduce the cache in memory
-        const tileSize = Math.max(tileWidth, tileHeight);
-        const tileOverlap = Math.max(tileOverlapX, tileOverlapY);
-        tileWidth = tileSize;
-        tileHeight = tileSize;
-        tileOverlapX = tileOverlap;
-        tileOverlapY = tileOverlap;
-        if (tileWidth !== this._tileWidth ||
-            tileHeight !== this._tileHeight ||
-            tileOverlapX !== this._tileOverlapX ||
-            tileOverlapY !== this._tileOverlapY) {
-            // console.log(tileWidth, tileHeight, tileOverlapX, tileOverlapY);
-            this._tileWidth = tileWidth;
-            this._tileHeight = tileHeight;
-            this._tileOverlapX = tileOverlapX;
-            this._tileOverlapY = tileOverlapY;
-        }
-    }
-    _getTileSizeWithOverlap() {
-        return {
-            width: this._tileWidth + 2 * this._tileOverlapX,
-            height: this._tileHeight + 2 * this._tileOverlapY
-        };
-    }
     _processImageData(color, albedo, normal, isHDR) {
         const rawData = color.data;
         const pixelsCount = rawData.length / 4;
@@ -179,9 +197,12 @@ class UNet {
     }
     _readTile(data, channels, srcTile, width) {
         const tileData = new Float32Array(srcTile.width * srcTile.height * channels);
+        const height = data.length / (width * channels);
         for (let y = 0; y < srcTile.height; y++) {
             for (let x = 0; x < srcTile.width; x++) {
-                const i2 = ((y + srcTile.y) * width + (x + srcTile.x)) * channels;
+                const sourceX = Math.min(width - 1, x + srcTile.x);
+                const sourceY = Math.min(height - 1, y + srcTile.y);
+                const i2 = (sourceY * width + sourceX) * channels;
                 const i1 = (y * srcTile.width + x) * channels;
                 for (let c = 0; c < channels; c++) {
                     tileData[i1 + c] = data[i2 + c];
@@ -210,21 +231,12 @@ class UNet {
             }
         }
     }
-    async _executeTile(inputData, outputTileData, outputImageData, i, j, width, height, isHDR, denoiseAlpha) {
+    async _executeTile(inputData, outputTileData, outputImageData, tile, isFirstTile, width, height, isHDR, denoiseAlpha) {
         const channels = this._aux ? 9 : 3;
-        const tileOverlapX = this._tileOverlapX;
-        const tileOverlapY = this._tileOverlapY;
-        let srcTileSize = this._getTileSizeWithOverlap();
-        let dstTileSize = { width: this._tileWidth, height: this._tileHeight };
-        let srcX0 = i > 0 ? i * dstTileSize.width - tileOverlapX : 0;
-        let srcX1 = Math.min(srcX0 + srcTileSize.width, width);
-        srcX0 = Math.max(srcX1 - srcTileSize.width, 0);
-        let srcY0 = j > 0 ? j * dstTileSize.height - tileOverlapY : 0;
-        let srcY1 = Math.min(srcY0 + srcTileSize.height, height);
-        srcY0 = Math.max(srcY1 - srcTileSize.height, 0);
-        const srcTileWidth = srcTileSize.width;
-        const srcTileHeight = srcTileSize.height;
-        const srcTile = new Tile(srcX0, srcY0, srcTileWidth, srcTileHeight);
+        const srcTile = new Tile(tile.input.x, tile.input.y, tile.input.width, tile.input.height);
+        const dstTile = new Tile(tile.output.x, tile.output.y, tile.output.width, tile.output.height);
+        const srcTileWidth = srcTile.width;
+        const srcTileHeight = srcTile.height;
         let nativeOutputBuffer;
         let denoisedData;
         let inputScale = 1;
@@ -240,42 +252,39 @@ class UNet {
                 tileData = hdrTransferFuncCPU({
                     data: tileData,
                     channels,
-                    inputScale
+                    inputScale,
+                    transfer: this._hdrTransfer
                 });
             }
             denoisedData = await (this._webNNExecutor ?? this._nativeExecutor).executeCPU(tileData, srcTileWidth, srcTileHeight);
         }
         else {
             if (!dataProcessGPU) {
-                dataProcessGPU = this._dataProcessGPU = new GPUDataProcess(device, isHDR);
+                dataProcessGPU = this._dataProcessGPU = new GPUDataProcess(device, isHDR, this._hdrTransfer);
             }
             dataProcessGPU.setImageSize(width, height);
             dataProcessGPU.setInputTile(srcTile);
             // Display the noisy input instead of prev denoised result
-            if (i === 0 && j === 0) {
+            if (isFirstTile) {
                 dataProcessGPU.copyInputDataToOutput(inputData.color);
             }
             const { color, albedo, normal } = dataProcessGPU.forward(inputData.color, this._aux ? inputData.albedo : undefined, this._aux ? inputData.normal : undefined, denoiseAlpha);
             nativeOutputBuffer = await (this._webNNExecutor ?? this._nativeExecutor).execute(this._aux ? [color, albedo, normal] : [color], srcTileWidth, srcTileHeight);
         }
         let outBuffer;
-        const dstWidth = Math.min(dstTileSize.width, width);
-        const dstHeight = Math.min(dstTileSize.height, height);
-        const dstTile = new Tile(i * dstWidth, j * dstHeight, dstWidth, dstHeight);
-        dstTile.width = Math.min(dstTile.width, width - dstTile.x);
-        dstTile.height = Math.min(dstTile.height, height - dstTile.y);
         if (inputData instanceof Float32Array) {
             if (isHDR) {
                 denoisedData = hdrTransferFuncInverseCPU({
                     data: denoisedData,
                     channels: 3,
-                    inputScale
+                    inputScale,
+                    transfer: this._hdrTransfer
                 });
             }
-            this._writeTile(outputImageData, srcTile, dstTile, denoisedData, srcTileSize.width, isHDR);
-            for (let y = 0; y < dstHeight; y++) {
-                for (let x = 0; x < dstWidth; x++) {
-                    const i1 = (y * dstWidth + x) * 4;
+            this._writeTile(outputImageData, srcTile, dstTile, denoisedData, srcTile.width, isHDR);
+            for (let y = 0; y < dstTile.height; y++) {
+                for (let x = 0; x < dstTile.width; x++) {
+                    const i1 = (y * dstTile.width + x) * 4;
                     const i2 = ((y + dstTile.y) * width + (x + dstTile.x)) * 4;
                     for (let c = 0; c < 4; c++) {
                         outputTileData.data[i1 + c] = outputImageData.data[i2 + c];
@@ -289,7 +298,7 @@ class UNet {
         }
         return outBuffer;
     }
-    tileExecute({ color, albedo, normal, done, progress, denoiseAlpha }) {
+    tileExecute({ color, albedo, normal, done, progress, denoiseAlpha, tileOverlap, wholeImage, scheduling = 'event-loop', error }) {
         if (this._aux && (!albedo || !normal)) {
             throw new Error('Normal map and albedo map are both required');
         }
@@ -301,8 +310,16 @@ class UNet {
         const width = color.width;
         const height = color.height;
         const adaptiveTileSize = this._dynamicTileController.tileSize;
-        const shouldAdaptTileSize = width > adaptiveTileSize || height > adaptiveTileSize;
-        this._updateModel(width, height);
+        // The planner aligns the maximum down, so round up to keep one tile.
+        const requestedTileSize = wholeImage
+            ? roundUp(Math.max(width, height), OIDN_TILE_ALIGNMENT)
+            : adaptiveTileSize;
+        const defaultTileOverlap = roundUp(this._modelSpec.receptiveField / 2, OIDN_TILE_ALIGNMENT);
+        const resolvedTileOverlap = tileOverlap === undefined
+            ? defaultTileOverlap
+            : roundUp(Math.max(0, tileOverlap), OIDN_TILE_ALIGNMENT);
+        const plan = planTileGrid(width, height, requestedTileSize, resolvedTileOverlap);
+        const shouldAdaptTileSize = plan.tiles.length > 1;
         // TODO should fixed to be hdr when UNet is created.
         // weights of hdr and ldr is different
         const hdr = this._hdr || false;
@@ -310,10 +327,6 @@ class UNet {
         if (!isGPUImageData(color)) {
             rawData = this._processImageData(color, albedo, normal, hdr);
         }
-        const tileWidth = this._tileWidth;
-        const tileHeight = this._tileHeight;
-        const tileCountH = Math.ceil(height / tileHeight);
-        const tileCountW = Math.ceil(width / tileWidth);
         function makeImageData(width, height) {
             return hdr
                 ? {
@@ -326,25 +339,75 @@ class UNet {
         const outputImageData = isGPUImageData(color)
             ? undefined
             : makeImageData(width, height);
-        const outputTileData = isGPUImageData(color)
-            ? undefined
-            : makeImageData(Math.min(tileWidth, width), Math.min(tileHeight, height));
-        let aborted = false;
+        let state = 'active';
+        let scheduledTimer;
+        let scheduledAnimationFrame;
+        let unregisterDeviceLoss = () => false;
         const now = () => typeof performance === 'undefined' ? Date.now() : performance.now();
         const executionStartTime = now();
         const tileTimesMs = [];
+        const cancelScheduledTile = () => {
+            if (scheduledTimer !== undefined) {
+                clearTimeout(scheduledTimer);
+                scheduledTimer = undefined;
+            }
+            if (scheduledAnimationFrame !== undefined &&
+                typeof cancelAnimationFrame !== 'undefined') {
+                cancelAnimationFrame(scheduledAnimationFrame);
+                scheduledAnimationFrame = undefined;
+            }
+        };
+        const reportCallbackFailure = (reason) => {
+            // An error callback is the terminal observer and cannot report its own
+            // failure through the same channel. Keep that failure handled.
+            console.error('OIDN error callback failed', reason);
+        };
+        const settleError = (reason) => {
+            if (state !== 'active')
+                return;
+            state = 'settled';
+            cancelScheduledTile();
+            unregisterDeviceLoss();
+            if (error) {
+                try {
+                    void Promise.resolve(error(reason)).catch(reportCallbackFailure);
+                }
+                catch (callbackReason) {
+                    reportCallbackFailure(callbackReason);
+                }
+            }
+            else {
+                console.error('OIDN execution failed', reason);
+            }
+        };
         const scheduleNextTile = (callback) => {
-            if (typeof requestAnimationFrame === 'undefined') {
-                setTimeout(callback, 0);
+            if (scheduling === 'event-loop' ||
+                typeof requestAnimationFrame === 'undefined') {
+                scheduledTimer = setTimeout(() => {
+                    scheduledTimer = undefined;
+                    callback();
+                }, 0);
             }
             else {
-                requestAnimationFrame(callback);
+                // Hidden documents pause requestAnimationFrame. Race it against a
+                // timer so animation-frame scheduling still completes in background
+                // tabs, minimized windows, and offscreen iframes.
+                const run = () => {
+                    cancelScheduledTile();
+                    callback();
+                };
+                scheduledAnimationFrame = requestAnimationFrame(run);
+                scheduledTimer = setTimeout(run, ANIMATION_FRAME_FALLBACK_MS);
             }
         };
-        const executeTile = async (i, j) => {
-            if (aborted) {
+        const executeTile = async (tileIndex) => {
+            if (state !== 'active') {
                 return;
             }
+            const tile = plan.tiles[tileIndex];
+            const outputTileData = isGPUImageData(color)
+                ? undefined
+                : makeImageData(tile.output.width, tile.output.height);
             const tileStartTime = now();
             const resGPUBuffer = await this._executeTile(isGPUImageData(color)
                 ? {
@@ -352,32 +415,34 @@ class UNet {
                     albedo: albedo?.data,
                     normal: normal?.data
                 }
-                : rawData, outputTileData, outputImageData, i, j, width, height, hdr, denoiseAlpha);
-            if (aborted)
+                : rawData, outputTileData, outputImageData, tile, tileIndex === 0, width, height, hdr, denoiseAlpha);
+            if (state !== 'active')
                 return;
             const output = outputImageData || {
                 data: resGPUBuffer,
                 width,
                 height
             };
-            progress?.(output, 
-            // Is undefined if using webgpu buffer
-            outputTileData, new Tile(i * tileWidth, j * tileHeight, tileWidth, tileHeight), i + j * tileCountW, tileCountW * tileCountH);
-            const hasNextTile = i + 1 < tileCountW || j + 1 < tileCountH;
-            const continueAfterGPUWork = () => {
+            if (progress) {
+                await progress(output, 
+                // Is undefined if using webgpu buffer
+                outputTileData, new Tile(tile.output.x, tile.output.y, tile.output.width, tile.output.height), tileIndex, plan.tiles.length);
+            }
+            if (state !== 'active')
+                return;
+            const hasNextTile = tileIndex + 1 < plan.tiles.length;
+            await this._device.queue.onSubmittedWorkDone();
+            if (state !== 'active')
+                return;
+            const continueAfterGPUWork = async () => {
                 tileTimesMs.push(now() - tileStartTime);
-                if (aborted)
+                if (state !== 'active')
                     return;
                 if (hasNextTile) {
                     scheduleNextTile(() => {
-                        if (aborted)
+                        if (state !== 'active')
                             return;
-                        if (i + 1 < tileCountW) {
-                            executeTile(i + 1, j);
-                        }
-                        else if (j + 1 < tileCountH) {
-                            executeTile(0, j + 1);
-                        }
+                        void executeTile(tileIndex + 1).catch(settleError);
                     });
                 }
                 else {
@@ -389,9 +454,12 @@ class UNet {
                     this._lastExecution = {
                         width,
                         height,
-                        tileWidth,
-                        tileHeight,
-                        tileCount: tileCountW * tileCountH,
+                        tileCount: plan.tiles.length,
+                        tileColumns: plan.columns,
+                        tileRows: plan.rows,
+                        tileOverlap: plan.overlap,
+                        inputPixelCount: plan.inputPixelCount,
+                        inputShapeCount: plan.inputShapeCount,
                         durationMs: now() - executionStartTime,
                         tileTimeMs: {
                             min: sortedTileTimes[0],
@@ -406,18 +474,26 @@ class UNet {
                     if (shouldAdaptTileSize) {
                         this._dynamicTileController.observe(tileTimesMs);
                     }
-                    // console.log(memory());
-                    done(output);
+                    // GPU inference is complete; device loss can no longer affect this
+                    // execution. Deregister before the user callback resolves so a
+                    // completed execution never remains retained by the device watcher.
+                    unregisterDeviceLoss();
+                    await done(output);
+                    if (state === 'active') {
+                        state = 'settled';
+                    }
                 }
             };
-            // requestAnimationFrame only throttles JavaScript submission. Waiting
-            // for the queue here keeps at most one OIDN tile in flight, so aborting
-            // cannot leave a long tail of already-submitted GPU work.
-            void waitForSubmittedGPUWork(this._device.queue).then(continueAfterGPUWork);
+            await continueAfterGPUWork();
         };
-        executeTile(0, 0);
+        unregisterDeviceLoss = this._registerExecutionFailure(settleError);
+        void executeTile(0).catch(settleError);
         return () => {
-            aborted = true;
+            if (state !== 'active')
+                return;
+            state = 'aborted';
+            cancelScheduledTile();
+            unregisterDeviceLoss();
         };
     }
     dispose() {