oidn-web 0.3.4 → 0.4.0

package/lib/UNet.d.ts

@@ -1,6 +1,9 @@
 import { HostTensor } from './tza';
 import { Tile } from './process';
-import type { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu';
+import { type DynamicTileSetting } from './tileScheduler';
+import { type UNetModelSpec } from './modelSpec';
+import { type NativeUNetKernelSetting, type NativeUNetPrecisionSetting } from './nativeUNet';
+export type UNetEngineSetting = 'auto' | 'wgsl' | 'webnn';
 interface HDRImageData {
     data: Float32Array;
     width: number;
@@ -16,10 +19,21 @@ interface GPUImageDataOutput {
     width: number;
     height: number;
 }
+export interface UNetExecutionStats {
+    width: number;
+    height: number;
+    tileWidth: number;
+    tileHeight: number;
+    tileCount: number;
+    durationMs: number;
+    tileTimeMs: {
+        min: number;
+        median: number;
+        mean: number;
+        max: number;
+    };
+}
 declare class UNet {
-    private _hostTensors;
-    private _backend;
-    private _tfModel;
     private _device;
     private _tileWidth;
     private _tileHeight;
@@ -28,10 +42,17 @@ declare class UNet {
     private _aux;
     private _hdr;
     private _dataProcessGPU?;
-    private _maxTileSize;
-    private _tensors;
-    private _modelsCache;
-    constructor(_hostTensors: Map<string, HostTensor>, _backend: WebGPUBackend, opts?: {
+    private _nativeExecutor?;
+    private _webNNExecutor?;
+    private _modelSpec;
+    private _inputChannels;
+    private _engine;
+    private _dynamicTileController;
+    private _lastExecution?;
+    constructor(hostTensors: Map<string, HostTensor>, backend: {
+        device: GPUDevice;
+        adapterInfo: GPUAdapterInfo;
+    }, opts?: {
         /**
          * If use auxiliary data.
          */
@@ -41,15 +62,45 @@ declare class UNet {
          */
         hdr?: boolean;
         maxTileSize?: number;
+        dynamicTile?: DynamicTileSetting;
+        /** Reserved for explicit native WGSL selection. */
+        engine?: UNetEngineSetting;
+        /** Arithmetic/storage precision used by the native WGSL engine. */
+        precision?: NativeUNetPrecisionSetting;
+        /** Model-independent convolution kernel selection. */
+        kernel?: NativeUNetKernelSetting;
+        /** Explicit descriptor for a new OIDN topology not in the built-in registry. */
+        modelSpec?: UNetModelSpec;
     });
     getDevice(): GPUDevice | undefined;
-    private _buildModel;
-    private _createConv;
-    private _createConcatConv;
-    private _createPooling;
-    private _addUpsamplingLayer;
-    private _addNet;
-    private _addNetLarge;
+    /** Completes backend compilation before first interactive use. */
+    prepare(): Promise<void>;
+    getRuntimeInfo(): {
+        configuredEngine: UNetEngineSetting;
+        gpuEngine: "wgsl" | "webnn";
+        precision: import("./nativeUNet").NativeUNetPrecision;
+        kernel: {
+            configured: NativeUNetKernelSetting;
+            maxSpatialInputBlocks: number;
+            subgroupsAvailable: boolean;
+        } | undefined;
+        webnn: import("./webnnUNet").WebNNRuntimeSupport | undefined;
+        resources: import("./resourceTracker").OIDNResourceSnapshot;
+        model: string;
+        modelFamily: (string & {}) | "oidn-unet-small" | "oidn-unet-large";
+        inputChannels: number;
+        dynamicTile: {
+            enabled: boolean;
+            currentTileSize: number;
+            minTileSize: number;
+            maxTileSize: number;
+            targetTileTimeMs: number;
+        };
+        lastExecution: UNetExecutionStats | undefined;
+    };
+    /** 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;

package/lib/UNet.js

@@ -1,67 +1,15 @@
-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, fitTileDimension, OIDN_TILE_ALIGNMENT, waitForSubmittedGPUWork } from './tileScheduler';
+import { detectUNetModelSpec, validateUNetModel } from './modelSpec';
+import { NativeUNetExecutor } from './nativeUNet';
+import { WebNNUNetExecutor } from './webnnUNet';
 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
@@ -72,164 +20,101 @@ class UNet {
     _aux;
     _hdr;
     _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;
+    constructor(hostTensors, backend, opts = {}) {
         this._aux = opts.aux || false;
         this._hdr = opts.hdr || false;
-        this._maxTileSize = roundUp(opts.maxTileSize ?? 512, 2);
-        this._device = this._backend.device;
+        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 = backend.device;
+        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 });
+        }
     }
     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);
+        await this._nativeExecutor.prepare();
     }
-    _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');
-    }
-    _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);
-    }
-    _addUpsamplingLayer(source) {
-        const upsamplingLayer = new UpSampling2D({
-            name: source.name + '/upsampling',
-            size: [2, 2],
-            trainable: false
-        });
-        return upsamplingLayer.apply(source);
+    getRuntimeInfo() {
+        return {
+            configuredEngine: this._engine,
+            gpuEngine: this._webNNExecutor ? 'webnn' : 'wgsl',
+            precision: (this._webNNExecutor ?? this._nativeExecutor).precision,
+            kernel: this._nativeExecutor
+                ? {
+                    configured: this._nativeExecutor.kernelSetting,
+                    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,
+            dynamicTile: {
+                enabled: this._dynamicTileController.enabled,
+                currentTileSize: this._dynamicTileController.tileSize,
+                minTileSize: this._dynamicTileController.minTileSize,
+                maxTileSize: this._dynamicTileController.maxTileSize,
+                targetTileTimeMs: this._dynamicTileController.targetTileTimeMs
+            },
+            lastExecution: this._lastExecution
+        };
     }
-    _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;
+    /** Captures per-node GPU timestamps for the next native tile execution. */
+    profileNextExecution() {
+        return this._nativeExecutor?.profileNextExecution() ?? false;
     }
-    _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;
+    getLastExecutionProfile() {
+        return this._nativeExecutor?.getLastExecutionProfile();
     }
     _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;
-            }
+        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 + defaultTileOverlap * 2) {
-            tileHeight = roundUp(height, maxTileSize / 2);
-            if (height <= maxTileSize) {
-                tileOverlapY = 0;
-            }
+        if (height <= maxTileSize) {
+            tileOverlapY = 0;
         }
         // Force width and height has same size. reduce the cache in memory
         const tileSize = Math.max(tileWidth, tileHeight);
@@ -241,14 +126,12 @@ class UNet {
         if (tileWidth !== this._tileWidth ||
             tileHeight !== this._tileHeight ||
             tileOverlapX !== this._tileOverlapX ||
-            tileOverlapY !== this._tileOverlapY ||
-            !this._tfModel) {
+            tileOverlapY !== this._tileOverlapY) {
             // console.log(tileWidth, tileHeight, tileOverlapX, tileOverlapY);
             this._tileWidth = tileWidth;
             this._tileHeight = tileHeight;
             this._tileOverlapX = tileOverlapX;
             this._tileOverlapY = tileOverlapY;
-            this._buildModel(isLarge);
         }
     }
     _getTileSizeWithOverlap() {
@@ -327,7 +210,7 @@ class UNet {
             }
         }
     }
-    _executeTile(inputData, outputTileData, outputImageData, i, j, width, height, isHDR, denoiseAlpha) {
+    async _executeTile(inputData, outputTileData, outputImageData, i, j, width, height, isHDR, denoiseAlpha) {
         const channels = this._aux ? 9 : 3;
         const tileOverlapX = this._tileOverlapX;
         const tileOverlapY = this._tileOverlapY;
@@ -342,7 +225,8 @@ class UNet {
         const srcTileWidth = srcTileSize.width;
         const srcTileHeight = srcTileSize.height;
         const srcTile = new Tile(srcX0, srcY0, srcTileWidth, srcTileHeight);
-        let tileTensor;
+        let nativeOutputBuffer;
+        let denoisedData;
         let inputScale = 1;
         const device = this._device;
         let dataProcessGPU = this._dataProcessGPU;
@@ -359,7 +243,7 @@ class UNet {
                     inputScale
                 });
             }
-            tileTensor = tensor(tileData, [1, srcTileHeight, srcTileWidth, channels], 'float32');
+            denoisedData = await (this._webNNExecutor ?? this._nativeExecutor).executeCPU(tileData, srcTileWidth, srcTileHeight);
         }
         else {
             if (!dataProcessGPU) {
@@ -372,33 +256,15 @@ class UNet {
                 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,
@@ -419,16 +285,7 @@ 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;
     }
@@ -443,6 +300,8 @@ 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);
         // TODO should fixed to be hdr when UNet is created.
         // weights of hdr and ldr is different
@@ -471,22 +330,31 @@ class UNet {
             ? undefined
             : makeImageData(Math.min(tileWidth, width), Math.min(tileHeight, height));
         let aborted = false;
-        const executeTile = (i, j) => {
+        const now = () => typeof performance === 'undefined' ? Date.now() : performance.now();
+        const executionStartTime = now();
+        const tileTimesMs = [];
+        const scheduleNextTile = (callback) => {
+            if (typeof requestAnimationFrame === 'undefined') {
+                setTimeout(callback, 0);
+            }
+            else {
+                requestAnimationFrame(callback);
+            }
+        };
+        const executeTile = async (i, j) => {
             if (aborted) {
                 return;
             }
-            let resGPUBuffer;
-            // profileAndLogKernelCode(() => {
-            ENGINE.startScope();
-            resGPUBuffer = this._executeTile(isGPUImageData(color)
+            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);
+            if (aborted)
+                return;
             const output = outputImageData || {
                 data: resGPUBuffer,
                 width,
@@ -495,20 +363,57 @@ class UNet {
             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);
-                    }
-                    else if (j + 1 < tileCountH) {
-                        executeTile(0, j + 1);
+            const hasNextTile = i + 1 < tileCountW || j + 1 < tileCountH;
+            const continueAfterGPUWork = () => {
+                tileTimesMs.push(now() - tileStartTime);
+                if (aborted)
+                    return;
+                if (hasNextTile) {
+                    scheduleNextTile(() => {
+                        if (aborted)
+                            return;
+                        if (i + 1 < tileCountW) {
+                            executeTile(i + 1, j);
+                        }
+                        else if (j + 1 < tileCountH) {
+                            executeTile(0, j + 1);
+                        }
+                    });
+                }
+                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,
+                        tileWidth,
+                        tileHeight,
+                        tileCount: tileCountW * tileCountH,
+                        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 {
-                // console.log(memory());
-                done(output);
-            }
+                    // console.log(memory());
+                    done(output);
+                }
+            };
+            // 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);
         };
         executeTile(0, 0);
         return () => {
@@ -516,9 +421,9 @@ class UNet {
         };
     }
     dispose() {
-        this._tfModel?.dispose();
         this._dataProcessGPU?.dispose();
-        this._tensors.forEach((tensor) => tensor.dispose());
+        this._nativeExecutor?.dispose();
+        this._webNNExecutor?.dispose();
     }
 }
 export default UNet;