oidn-web 0.3.5 → 0.5.0

package/lib/UNet.js

@@ -1,261 +1,162 @@
-import { tensor } from '@tensorflow/tfjs-core/dist/ops/tensor';
-import { tensor1d } from '@tensorflow/tfjs-core/dist/ops/tensor1d';
-import { pad4d } from '@tensorflow/tfjs-core/dist/ops/pad4d';
-import { slice4d } from '@tensorflow/tfjs-core/dist/ops/slice4d';
-import { concat4d } from '@tensorflow/tfjs-core/dist/ops/concat_4d';
-import { ENGINE } from '@tensorflow/tfjs-core/dist/engine';
-import { Conv2D, UpSampling2D } from '@tensorflow/tfjs-layers/dist/layers/convolutional';
-import { MaxPooling2D } from '@tensorflow/tfjs-layers/dist/layers/pooling';
-import { Concatenate } from '@tensorflow/tfjs-layers/dist/layers/merge';
-import { LayersModel } from '@tensorflow/tfjs-layers/dist/engine/training';
-import { Input as TFInput } from '@tensorflow/tfjs-layers/dist/engine/input_layer';
-import { Float16Array } from '@petamoriken/float16';
 import { GPUDataProcess, Tile, avgLogLum, hdrTransferFuncCPU, hdrTransferFuncInverseCPU } from './process';
-// import { profileAndLogKernelCode, memory } from './helper';
-function getTensorData(ubytes, type) {
-    const buffer = ubytes.buffer;
-    if (type === 'Float32') {
-        return new Float32Array(ubytes.buffer);
-    }
-    const float16Data = new Float16Array(buffer);
-    const float32Data = new Float32Array(float16Data.length);
-    for (let i = 0; i < float32Data.length; ++i) {
-        float32Data[i] = float16Data[i];
-    }
-    return float32Data;
-}
-function changeWeightShapes(weightData, dims) {
-    const [O, C, H, W] = dims;
-    const reorderedWeightData = new Float32Array(weightData.length);
-    for (let o = 0; o < O; ++o) {
-        for (let c = 0; c < C; ++c) {
-            for (let h = 0; h < H; ++h) {
-                for (let w = 0; w < W; ++w) {
-                    // Change OCHW to HWCO
-                    const idx = o * C * H * W + c * H * W + h * W + w;
-                    const idx2 = h * W * C * O + w * C * O + c * O + o;
-                    reorderedWeightData[idx2] = weightData[idx];
-                }
-            }
-        }
-    }
-    return reorderedWeightData;
-}
+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;
 }
-// Returns the smallest integer larger than or equal to a which has remainder c when divided by b
-function roundUp2(a, b, c) {
-    return Math.ceil((a - c) / b) * b + c;
-}
 function isGPUImageData(data) {
     return data.data instanceof GPUBuffer || data.data instanceof GPUTexture;
 }
-const receptiveField = 174; // receptive field in pixels
-const receptiveFieldLarge = 202;
-// TODO metal is 32?
-const minTileAlignment = 1;
-const tileAlignment = 16; // required spatial alignment in pixels (padding may be necessary)
-const defaultTileOverlap = roundUp(receptiveField / 2, tileAlignment);
-const defaultTileOverlapLarge = roundUp(receptiveFieldLarge / 2, tileAlignment);
 class UNet {
-    _hostTensors;
-    _backend;
-    _tfModel;
     _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;
-    _maxTileSize;
-    _tensors = new Map();
-    _modelsCache = new Map();
-    constructor(_hostTensors, _backend, opts = {}) {
-        this._hostTensors = _hostTensors;
-        this._backend = _backend;
+    _nativeExecutor;
+    _webNNExecutor;
+    _modelSpec;
+    _inputChannels;
+    _engine;
+    _dynamicTileController;
+    _lastExecution;
+    _activeExecutionFailures = new Set();
+    _deviceLostObserved = false;
+    _deviceLostSettled = false;
+    _deviceLostReason;
+    constructor(hostTensors, device, opts = {}) {
         this._aux = opts.aux || false;
         this._hdr = opts.hdr || false;
-        this._maxTileSize = roundUp(opts.maxTileSize ?? 512, 2);
-        this._device = this._backend.device;
+        this._hdrTransfer = opts.hdrTransfer ?? 'pu';
+        this._engine = opts.engine ?? 'auto';
+        const modelSpec = opts.modelSpec ?? detectUNetModelSpec(hostTensors);
+        const validatedModel = validateUNetModel(hostTensors, modelSpec);
+        this._modelSpec = validatedModel.spec;
+        this._inputChannels = validatedModel.inputChannels;
+        const expectedInputChannels = this._aux ? 9 : 3;
+        if (validatedModel.inputChannels !== expectedInputChannels) {
+            throw new Error(`OIDN model expects ${validatedModel.inputChannels} input channels, ` +
+                `but aux=${this._aux} provides ${expectedInputChannels}`);
+        }
+        this._dynamicTileController = new DynamicTileController(opts.maxTileSize ?? 512, opts.dynamicTile);
+        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, gemm: opts.gemm });
+        }
     }
     getDevice() {
         return this._device;
     }
-    _buildModel(isLarge) {
-        const aux = this._aux;
-        const channels = 3 + (aux ? 6 : 0);
-        const tileSize = this._getTileSizeWithOverlap();
-        const cache = this._modelsCache;
-        const key = [tileSize.width, tileSize.height].join(',');
-        // We cache the model instead of disposing and recreate.
-        // Because seems tfjs will also cache the layer and gpubuffers.
-        // Recreating the model will cause memory leak.
-        // Width and height can only be 256, 512, 768. So the cache won't be too large
-        if (cache.has(key)) {
-            this._tfModel = cache.get(key);
+    /** Completes backend compilation before first interactive use. */
+    async prepare() {
+        if (this._webNNExecutor) {
+            await this._webNNExecutor.prepare();
+            const overlap = roundUp(this._modelSpec.receptiveField / 2, OIDN_TILE_ALIGNMENT);
+            const outputTileEdges = [
+                this._dynamicTileController.tileSize,
+                this._dynamicTileController.minTileSize
+            ];
+            await this._webNNExecutor.prewarm([...new Set(outputTileEdges)].map((edge) => ({
+                width: edge + 2 * overlap,
+                height: edge + 2 * overlap
+            })));
             return;
         }
-        const input = TFInput({
-            name: 'input',
-            shape: [tileSize.height, tileSize.width, channels],
-            dtype: 'float32'
-        });
-        this._tfModel = new LayersModel({
-            inputs: [input],
-            outputs: isLarge ? this._addNetLarge(input) : this._addNet(input)
-        });
-        cache.set(key, this._tfModel);
-    }
-    _createConv(name, source, activation) {
-        const weightTensorName = name + '.weight';
-        const biasTensorName = name + '.bias';
-        const tensors = this._tensors;
-        let weightTensor = tensors.get(weightTensorName);
-        let biasTensor = tensors.get(biasTensorName);
-        const unetWeightTensor = this._hostTensors.get(weightTensorName);
-        if (!weightTensor) {
-            const weightDims = unetWeightTensor.desc.dims;
-            weightTensor = tensor(changeWeightShapes(getTensorData(unetWeightTensor.data, unetWeightTensor.desc.dataType), weightDims), [weightDims[2], weightDims[3], weightDims[1], weightDims[0]], 'float32');
-            tensors.set(weightTensorName, weightTensor);
-        }
-        if (!biasTensor) {
-            const unetBiasTensor = this._hostTensors.get(name + '.bias');
-            biasTensor = tensor1d(getTensorData(unetBiasTensor.data, unetBiasTensor.desc.dataType), 'float32');
-            tensors.set(biasTensorName, biasTensor);
-        }
-        // TODO whats the purpose of padded dims ?
-        const convLayer = new Conv2D({
-            name,
-            filters: unetWeightTensor.desc.dims[0],
-            kernelSize: unetWeightTensor.desc.dims.slice(2, 4),
-            useBias: true,
-            activation,
-            padding: 'same',
-            weights: [weightTensor, biasTensor],
-            trainable: false
-        });
-        return convLayer.apply(source);
-    }
-    _createConcatConv(name, source1, source2) {
-        const concatLayer = new Concatenate({
-            name: name + '/concat',
-            trainable: false,
-            axis: 3
-        });
-        //https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/training/model.py#L40
-        return this._createConv(name, 
-        // Concat on the channel
-        concatLayer.apply([source1, source2]), 'relu');
+        await this._nativeExecutor.prepare();
     }
-    _createPooling(source) {
-        const poolingLayer = new MaxPooling2D({
-            name: source.name + '/pooling',
-            poolSize: [2, 2],
-            strides: [2, 2],
-            padding: 'same',
-            trainable: false
-        });
-        // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/training/model.py#L33
-        return poolingLayer.apply(source);
+    /**
+     * 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);
     }
-    _addUpsamplingLayer(source) {
-        const upsamplingLayer = new UpSampling2D({
-            name: source.name + '/upsampling',
-            size: [2, 2],
-            trainable: false
-        });
-        return upsamplingLayer.apply(source);
-    }
-    _addNet(input) {
-        let x = this._createConv('enc_conv0', input, 'relu');
-        const pool1 = (x = this._createPooling(this._createConv('enc_conv1', x, 'relu')));
-        const pool2 = (x = this._createPooling(this._createConv('enc_conv2', x, 'relu')));
-        const pool3 = (x = this._createPooling(this._createConv('enc_conv3', x, 'relu')));
-        const pool4 = (x = this._createPooling(this._createConv('enc_conv4', x, 'relu')));
-        x = this._createConv('enc_conv5a', pool4, 'relu');
-        x = this._addUpsamplingLayer(this._createConv('enc_conv5b', x, 'relu'));
-        x = this._createConcatConv('dec_conv4a', x, pool3);
-        x = this._addUpsamplingLayer(this._createConv('dec_conv4b', x, 'relu'));
-        x = this._createConcatConv('dec_conv3a', x, pool2);
-        x = this._addUpsamplingLayer(this._createConv('dec_conv3b', x, 'relu'));
-        x = this._createConcatConv('dec_conv2a', x, pool1);
-        x = this._addUpsamplingLayer(this._createConv('dec_conv2b', x, 'relu'));
-        x = this._createConcatConv('dec_conv1a', x, input);
-        x = this._createConv('dec_conv1b', x, 'relu');
-        x = this._createConv('dec_conv0', x, 'relu');
-        return x;
+    getRuntimeInfo() {
+        return {
+            configuredEngine: this._engine,
+            gpuEngine: this._webNNExecutor ? 'webnn' : 'wgsl',
+            precision: (this._webNNExecutor ?? this._nativeExecutor).precision,
+            kernel: this._nativeExecutor
+                ? {
+                    configured: this._nativeExecutor.kernelSetting,
+                    gemm: this._nativeExecutor.gemm,
+                    maxSpatialInputBlocks: this._nativeExecutor.maxSpatialInputBlocks,
+                    subgroupsAvailable: this._nativeExecutor.subgroupsAvailable
+                }
+                : undefined,
+            webnn: this._webNNExecutor?.support,
+            resources: (this._webNNExecutor ?? this._nativeExecutor).getResourceInfo(),
+            model: this._modelSpec.id,
+            modelFamily: this._modelSpec.family,
+            inputChannels: this._inputChannels,
+            hdrTransfer: this._hdrTransfer,
+            dynamicTile: {
+                enabled: this._dynamicTileController.enabled,
+                currentTileSize: this._dynamicTileController.tileSize,
+                minTileSize: this._dynamicTileController.minTileSize,
+                maxTileSize: this._dynamicTileController.maxTileSize,
+                targetTileTimeMs: this._dynamicTileController.targetTileTimeMs
+            },
+            lastExecution: this._lastExecution,
+            activeExecutionCount: this._activeExecutionFailures.size
+        };
     }
-    _addNetLarge(input) {
-        let x = this._createConv('enc_conv1a', input, 'relu');
-        const pool1 = (x = this._createPooling(this._createConv('enc_conv1b', x, 'relu')));
-        x = this._createConv('enc_conv2a', x, 'relu');
-        const pool2 = (x = this._createPooling(this._createConv('enc_conv2b', x, 'relu')));
-        x = this._createConv('enc_conv3a', x, 'relu');
-        const pool3 = (x = this._createPooling(this._createConv('enc_conv3b', x, 'relu')));
-        x = this._createConv('enc_conv4a', x, 'relu');
-        const pool4 = (x = this._createPooling(this._createConv('enc_conv4b', x, 'relu')));
-        x = this._createConv('enc_conv5a', pool4, 'relu');
-        x = this._addUpsamplingLayer(this._createConv('enc_conv5b', x, 'relu'));
-        x = this._createConcatConv('dec_conv4a', x, pool3);
-        x = this._addUpsamplingLayer(this._createConv('dec_conv4b', x, 'relu'));
-        x = this._createConcatConv('dec_conv3a', x, pool2);
-        x = this._addUpsamplingLayer(this._createConv('dec_conv3b', x, 'relu'));
-        x = this._createConcatConv('dec_conv2a', x, pool1);
-        x = this._addUpsamplingLayer(this._createConv('dec_conv2b', x, 'relu'));
-        x = this._createConcatConv('dec_conv1a', x, input);
-        x = this._createConv('dec_conv1b', x, 'relu');
-        x = this._createConv('dec_conv1c', x, 'relu');
-        return x;
+    _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);
     }
-    _updateModel(width, height) {
-        const isLarge = this._hostTensors.has('enc_conv1b.weight');
-        const maxTileSize = this._maxTileSize;
-        let tileWidth = maxTileSize;
-        let tileHeight = maxTileSize;
-        let tileOverlapX = isLarge ? defaultTileOverlapLarge : defaultTileOverlap;
-        let tileOverlapY = isLarge ? defaultTileOverlapLarge : defaultTileOverlap;
-        if (width < maxTileSize + defaultTileOverlap * 2) {
-            tileWidth = roundUp(width, maxTileSize / 2);
-            if (width <= maxTileSize) {
-                tileOverlapX = 0;
-            }
-        }
-        if (height < maxTileSize + defaultTileOverlap * 2) {
-            tileHeight = roundUp(height, maxTileSize / 2);
-            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 ||
-            !this._tfModel) {
-            // console.log(tileWidth, tileHeight, tileOverlapX, tileOverlapY);
-            this._tileWidth = tileWidth;
-            this._tileHeight = tileHeight;
-            this._tileOverlapX = tileOverlapX;
-            this._tileOverlapY = tileOverlapY;
-            this._buildModel(isLarge);
+    _registerExecutionFailure(fail) {
+        this._observeDeviceLoss();
+        this._activeExecutionFailures.add(fail);
+        if (this._deviceLostSettled) {
+            this._activeExecutionFailures.delete(fail);
+            queueMicrotask(() => fail(this._deviceLostReason));
         }
+        return () => this._activeExecutionFailures.delete(fail);
     }
-    _getTileSizeWithOverlap() {
-        return {
-            width: this._tileWidth + 2 * this._tileOverlapX,
-            height: this._tileHeight + 2 * this._tileOverlapY
-        };
+    /** Captures per-node GPU timestamps for the next native tile execution. */
+    profileNextExecution() {
+        return this._nativeExecutor?.profileNextExecution() ?? false;
+    }
+    getLastExecutionProfile() {
+        return this._nativeExecutor?.getLastExecutionProfile();
     }
     _processImageData(color, albedo, normal, isHDR) {
         const rawData = color.data;
@@ -296,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];
@@ -327,22 +231,14 @@ class UNet {
             }
         }
     }
-    _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);
-        let tileTensor;
+        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;
         const device = this._device;
         let dataProcessGPU = this._dataProcessGPU;
@@ -356,60 +252,39 @@ class UNet {
                 tileData = hdrTransferFuncCPU({
                     data: tileData,
                     channels,
-                    inputScale
+                    inputScale,
+                    transfer: this._hdrTransfer
                 });
             }
-            tileTensor = tensor(tileData, [1, srcTileHeight, srcTileWidth, channels], 'float32');
+            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);
-            const createTensor = (buffer) => {
-                const tmp = tensor({ buffer, zeroCopy: true }, [
-                    1,
-                    srcTileHeight,
-                    srcTileWidth,
-                    4
-                ]);
-                const ret = slice4d(tmp, [0, 0, 0, 0], [1, srcTileHeight, srcTileWidth, 3]);
-                return ret;
-            };
-            if (this._aux) {
-                const tensors = [color, albedo, normal].map((buffer) => createTensor(buffer));
-                tileTensor = concat4d(tensors, 3);
-            }
-            else {
-                tileTensor = createTensor(color);
-            }
+            nativeOutputBuffer = await (this._webNNExecutor ?? this._nativeExecutor).execute(this._aux ? [color, albedo, normal] : [color], srcTileWidth, srcTileHeight);
         }
         let outBuffer;
-        const outputTensor = this._tfModel.predict(tileTensor);
-        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) {
-            let denoisedData = outputTensor.dataSync();
             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];
@@ -419,20 +294,11 @@ class UNet {
         }
         else {
             dataProcessGPU.setOutputTile(dstTile, srcTile);
-            // IMPORTANT
-            // storage buffer has alignment. that 3 channels still needs 16 bytes data.
-            // So we need to pad it to 4 channels.
-            const outputTensor4Channnels = pad4d(outputTensor, [
-                [0, 0],
-                [0, 0],
-                [0, 0],
-                [0, 1]
-            ]);
-            outBuffer = dataProcessGPU.inverse(outputTensor4Channnels.dataToGPU().buffer, inputData.color);
+            outBuffer = dataProcessGPU.inverse(nativeOutputBuffer, inputData.color);
         }
         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');
         }
@@ -443,7 +309,17 @@ class UNet {
         }
         const width = color.width;
         const height = color.height;
-        this._updateModel(width, height);
+        const adaptiveTileSize = this._dynamicTileController.tileSize;
+        // 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;
@@ -451,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
                 ? {
@@ -467,58 +339,167 @@ 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;
-        const executeTile = (i, j) => {
-            if (aborted) {
+        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 (scheduling === 'event-loop' ||
+                typeof requestAnimationFrame === 'undefined') {
+                scheduledTimer = setTimeout(() => {
+                    scheduledTimer = undefined;
+                    callback();
+                }, 0);
+            }
+            else {
+                // 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);
             }
-            let resGPUBuffer;
-            // profileAndLogKernelCode(() => {
-            ENGINE.startScope();
-            resGPUBuffer = this._executeTile(isGPUImageData(color)
+        };
+        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)
                 ? {
                     color: color.data,
                     albedo: albedo?.data,
                     normal: normal?.data
                 }
-                : rawData, outputTileData, outputImageData, i, j, width, height, hdr, denoiseAlpha);
-            ENGINE.endScope();
-            // }, true);
+                : 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);
-            if (i + 1 < tileCountW || j + 1 < tileCountH) {
-                requestAnimationFrame(() => {
-                    if (i + 1 < tileCountW) {
-                        executeTile(i + 1, j);
+            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 (state !== 'active')
+                    return;
+                if (hasNextTile) {
+                    scheduleNextTile(() => {
+                        if (state !== 'active')
+                            return;
+                        void executeTile(tileIndex + 1).catch(settleError);
+                    });
+                }
+                else {
+                    const sortedTileTimes = [...tileTimesMs].sort((a, b) => a - b);
+                    const middle = Math.floor(sortedTileTimes.length / 2);
+                    const medianTileTime = sortedTileTimes.length % 2
+                        ? sortedTileTimes[middle]
+                        : (sortedTileTimes[middle - 1] + sortedTileTimes[middle]) / 2;
+                    this._lastExecution = {
+                        width,
+                        height,
+                        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],
+                            median: medianTileTime,
+                            mean: sortedTileTimes.reduce((sum, value) => sum + value, 0) /
+                                sortedTileTimes.length,
+                            max: sortedTileTimes[sortedTileTimes.length - 1]
+                        }
+                    };
+                    // Adapt only from complete executions. Cancelled work is commonly
+                    // contending with interactive rendering and is not representative.
+                    if (shouldAdaptTileSize) {
+                        this._dynamicTileController.observe(tileTimesMs);
                     }
-                    else if (j + 1 < tileCountH) {
-                        executeTile(0, j + 1);
+                    // 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';
                     }
-                });
-            }
-            else {
-                // console.log(memory());
-                done(output);
-            }
+                }
+            };
+            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() {
-        this._tfModel?.dispose();
         this._dataProcessGPU?.dispose();
-        this._tensors.forEach((tensor) => tensor.dispose());
+        this._nativeExecutor?.dispose();
+        this._webNNExecutor?.dispose();
     }
 }
 export default UNet;