oidn-web 0.3.5 → 0.5.0

package/src/UNet.ts

@@ -1,67 +1,36 @@
-// import * as tfjs from '@tensorflow/tfjs-core';
-import { Tensor, Tensor1D, Tensor4D } from '@tensorflow/tfjs-core';
-import type { SymbolicTensor } from '@tensorflow/tfjs-layers';
-import { tensor } from '@tensorflow/tfjs-core/dist/ops/tensor';
-import { tensor1d } from '@tensorflow/tfjs-core/dist/ops/tensor1d';
-import { mirrorPad } from '@tensorflow/tfjs-core/dist/ops/mirror_pad';
-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 { HostTensor } from './tza';
-import { Float16Array } from '@petamoriken/float16';
 import {
   GPUDataProcess,
   Tile,
   avgLogLum,
   hdrTransferFuncCPU,
-  hdrTransferFuncInverseCPU
+  hdrTransferFuncInverseCPU,
+  type HDRTransfer
 } from './process';
-import type { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu';
+import {
+  DynamicTileController,
+  type DynamicTileSetting,
+  planTileGrid,
+  type PlannedTile,
+  OIDN_TILE_ALIGNMENT
+} from './tileScheduler';
+import {
+  detectUNetModelSpec,
+  validateUNetModel,
+  type UNetModelSpec
+} from './modelSpec';
+import {
+  NativeUNetExecutor,
+  type NativeUNetGemmOptions,
+  type NativeUNetKernelSetting,
+  type NativeUNetPrecisionSetting
+} from './nativeUNet';
+import { WebNNUNetExecutor } from './webnnUNet';
 
-// import { profileAndLogKernelCode, memory } from './helper';
+export type UNetEngineSetting = 'auto' | 'wgsl' | 'webnn';
 
-function getTensorData(
-  ubytes: Uint8Array,
-  type: HostTensor['desc']['dataType']
-) {
-  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: Float32Array, dims: number[]) {
-  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;
-}
+/** Upper bound on waiting for a display frame between tiles. */
+const ANIMATION_FRAME_FALLBACK_MS = 100;
 
 interface HDRImageData {
   data: Float32Array;
@@ -81,13 +50,27 @@ interface GPUImageDataOutput {
   height: number;
 }
 
+export interface UNetExecutionStats {
+  width: number;
+  height: number;
+  tileCount: number;
+  tileColumns: number;
+  tileRows: number;
+  tileOverlap: number;
+  inputPixelCount: number;
+  inputShapeCount: number;
+  durationMs: number;
+  tileTimeMs: {
+    min: number;
+    median: number;
+    mean: number;
+    max: number;
+  };
+}
+
 function roundUp(a: number, b: number) {
   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: number, b: number, c: number) {
-  return Math.ceil((a - c) / b) * b + c;
-}
 
 function isGPUImageData(
   data: ImageData | GPUImageData | HDRImageData
@@ -95,41 +78,30 @@ function isGPUImageData(
   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 {
-  private _tfModel: LayersModel | undefined;
-  private _device: GPUDevice | undefined;
-
-  // TODO calculate the tile size from memory size
-  // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/core/unet_filter.cpp#L287
-  private _tileWidth = 0;
-  private _tileHeight = 0;
-
-  private _tileOverlapX = 0;
-  private _tileOverlapY = 0;
+  private _device: GPUDevice;
 
   private _aux;
   private _hdr;
+  private _hdrTransfer: HDRTransfer;
 
   private _dataProcessGPU?: GPUDataProcess;
-
-  private _maxTileSize;
-
-  private _tensors = new Map<string, Tensor>();
-  private _modelsCache = new Map<string, LayersModel>();
+  private _nativeExecutor?: NativeUNetExecutor;
+  private _webNNExecutor?: WebNNUNetExecutor;
+  private _modelSpec: UNetModelSpec;
+  private _inputChannels: number;
+  private _engine: UNetEngineSetting;
+
+  private _dynamicTileController: DynamicTileController;
+  private _lastExecution?: UNetExecutionStats;
+  private _activeExecutionFailures = new Set<(reason: unknown) => void>();
+  private _deviceLostObserved = false;
+  private _deviceLostSettled = false;
+  private _deviceLostReason: unknown;
 
   constructor(
-    private _hostTensors: Map<string, HostTensor>,
-    private _backend: WebGPUBackend,
+    hostTensors: Map<string, HostTensor>,
+    device: GPUDevice,
     opts: {
       /**
        * If use auxiliary data.
@@ -139,261 +111,194 @@ class UNet {
        * If input is HDR image.
        */
       hdr?: boolean;
+      /** HDR transfer function expected by the trained model. */
+      hdrTransfer?: HDRTransfer;
       maxTileSize?: number;
+      dynamicTile?: DynamicTileSetting;
+      /** Native WGSL or the experimental WebNN backend. */
+      engine?: UNetEngineSetting;
+      /** 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;
     } = {}
   ) {
     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);
+    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._maxTileSize = roundUp(opts.maxTileSize ?? 512, 2);
+    this._dynamicTileController = new DynamicTileController(
+      opts.maxTileSize ?? 512,
+      opts.dynamicTile
+    );
 
-    this._device = this._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, gemm: opts.gemm }
+      );
+    }
   }
 
   getDevice() {
     return this._device;
   }
 
-  private _buildModel(isLarge: boolean) {
-    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);
-      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);
-  }
-
-  private _createConv(
-    name: string,
-    source: SymbolicTensor,
-    activation?: 'relu'
-  ) {
-    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'
+  /** 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
       );
-
-      tensors.set(weightTensorName, weightTensor);
-    }
-    if (!biasTensor) {
-      const unetBiasTensor = this._hostTensors.get(name + '.bias')!;
-      biasTensor = tensor1d(
-        getTensorData(unetBiasTensor.data, unetBiasTensor.desc.dataType),
-        'float32'
+      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
+        }))
       );
-      tensors.set(biasTensorName, biasTensor);
+      return;
     }
-    // TODO whats the purpose of padded dims ?
-    const convLayer = new Conv2D({
-      name,
-      filters: unetWeightTensor.desc.dims[0],
-      kernelSize: unetWeightTensor.desc.dims.slice(2, 4) as [number, number],
-      useBias: true,
-      activation,
-      padding: 'same',
-      weights: [weightTensor, biasTensor],
-      trainable: false
-    });
-
-    return convLayer.apply(source) as SymbolicTensor;
+    await this._nativeExecutor!.prepare();
   }
 
-  private _createConcatConv(
-    name: string,
-    source1: SymbolicTensor,
-    source2: SymbolicTensor
+  /**
+   * 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: number,
+    height: number,
+    options: { tileOverlap?: number; wholeImage?: boolean } = {}
   ) {
-    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]) as SymbolicTensor,
-      'relu'
-    ) as SymbolicTensor;
-  }
-
-  private _createPooling(source: SymbolicTensor) {
-    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) as SymbolicTensor;
-  }
-
-  private _addUpsamplingLayer(source: SymbolicTensor) {
-    const upsamplingLayer = new UpSampling2D({
-      name: source.name + '/upsampling',
-      size: [2, 2],
-      trainable: false
-    });
-    return upsamplingLayer.apply(source) as SymbolicTensor;
+    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);
   }
 
-  private _addNet(input: SymbolicTensor) {
-    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' as const : 'wgsl' as const,
+      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
+    };
   }
 
-  private _addNetLarge(input: SymbolicTensor) {
-    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;
+  private _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 as GPUDevice & {
+      lost?: Promise<GPUDeviceLostInfo>;
+    }).lost;
+    if (!deviceLost) return;
+    const failAll = (reason: unknown) => {
+      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
+    );
   }
 
-  private _updateModel(width: number, height: number) {
-    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;
-      }
+  private _registerExecutionFailure(fail: (reason: unknown) => void) {
+    this._observeDeviceLoss();
+    this._activeExecutionFailures.add(fail);
+    if (this._deviceLostSettled) {
+      this._activeExecutionFailures.delete(fail);
+      queueMicrotask(() => fail(this._deviceLostReason));
     }
+    return () => this._activeExecutionFailures.delete(fail);
+  }
 
-    // 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);
-    }
+  /** Captures per-node GPU timestamps for the next native tile execution. */
+  profileNextExecution() {
+    return this._nativeExecutor?.profileNextExecution() ?? false;
   }
 
-  private _getTileSizeWithOverlap() {
-    return {
-      width: this._tileWidth + 2 * this._tileOverlapX,
-      height: this._tileHeight + 2 * this._tileOverlapY
-    };
+  getLastExecutionProfile() {
+    return this._nativeExecutor?.getLastExecutionProfile();
   }
 
   private _processImageData(
@@ -453,9 +358,12 @@ class UNet {
     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++) {
@@ -497,7 +405,7 @@ class UNet {
     }
   }
 
-  private _executeTile(
+  private async _executeTile(
     inputData:
       | Float32Array
       | {
@@ -508,35 +416,33 @@ class UNet {
         },
     outputTileData: ImageData | HDRImageData | undefined,
     outputImageData: ImageData | HDRImageData | undefined,
-    i: number,
-    j: number,
+    tile: PlannedTile,
+    isFirstTile: boolean,
     width: number,
     height: number,
     isHDR: boolean,
     denoiseAlpha?: boolean
   ) {
     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 tileTensor!: Tensor;
+    let nativeOutputBuffer: GPUBuffer | undefined;
+    let denoisedData: Float32Array | undefined;
     let inputScale = 1;
-    const device = this._device!;
+    const device = this._device;
     let dataProcessGPU = this._dataProcessGPU;
 
     if (inputData instanceof Float32Array) {
@@ -549,25 +455,27 @@ class UNet {
         tileData = hdrTransferFuncCPU({
           data: tileData,
           channels,
-          inputScale
+          inputScale,
+          transfer: this._hdrTransfer
         });
       }
-      tileTensor = tensor(
+      denoisedData = await (this._webNNExecutor ?? this._nativeExecutor!).executeCPU(
         tileData,
-        [1, srcTileHeight, srcTileWidth, channels],
-        'float32'
-      ) as Tensor4D;
+        srcTileWidth,
+        srcTileHeight
+      );
     } else {
       if (!dataProcessGPU) {
         dataProcessGPU = this._dataProcessGPU = new GPUDataProcess(
           device,
-          isHDR
+          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(
@@ -577,47 +485,22 @@ class UNet {
         denoiseAlpha
       );
 
-      const createTensor = (buffer: GPUBuffer) => {
-        const tmp = tensor({ buffer, zeroCopy: true }, [
-          1,
-          srcTileHeight,
-          srcTileWidth,
-          4
-        ]) as Tensor4D;
-        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: GPUBuffer;
-    const outputTensor = this._tfModel!.predict(tileTensor) as Tensor;
-
-    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 as Float32Array,
+          data: denoisedData!,
           channels: 3,
-          inputScale
+          inputScale,
+          transfer: this._hdrTransfer
         });
       }
 
@@ -625,14 +508,14 @@ class UNet {
         outputImageData!,
         srcTile,
         dstTile,
-        denoisedData as Float32Array,
-        srcTileSize.width,
+        denoisedData!,
+        srcTile.width,
         isHDR
       );
 
-      for (let y = 0; y < dstHeight; y++) {
-        for (let x = 0; x < dstWidth; x++) {
-          const i1 = (y * dstWidth + x) * 4;
+      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];
@@ -641,17 +524,8 @@ 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 as Tensor4D, [
-        [0, 0],
-        [0, 0],
-        [0, 0],
-        [0, 1]
-      ]);
       outBuffer = dataProcessGPU!.inverse(
-        outputTensor4Channnels.dataToGPU().buffer!,
+        nativeOutputBuffer!,
         inputData.color
       );
     }
@@ -664,7 +538,11 @@ class UNet {
     normal,
     done,
     progress,
-    denoiseAlpha
+    denoiseAlpha,
+    tileOverlap,
+    wholeImage,
+    scheduling = 'event-loop',
+    error
   }: {
     color: T;
     albedo?: ImageData | GPUImageData;
@@ -673,14 +551,35 @@ class UNet {
      * If denoise alpha channel. Otherwise denoise RGB channels.
      */
     denoiseAlpha?: boolean;
-    done: (outputData: T extends GPUImageData ? GPUImageDataOutput : T) => 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;
+    ) => void | Promise<void>;
   }): () => void {
     if (this._aux && (!albedo || !normal)) {
       throw new Error('Normal map and albedo map are both required');
@@ -694,7 +593,25 @@ 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
@@ -709,11 +626,6 @@ class UNet {
         hdr
       );
     }
-    const tileWidth = this._tileWidth;
-    const tileHeight = this._tileHeight;
-    const tileCountH = Math.ceil(height / tileHeight);
-    const tileCountW = Math.ceil(width / tileWidth);
-
     function makeImageData(width: number, height: number) {
       return hdr
         ? {
@@ -727,20 +639,82 @@ 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;
+    type ExecutionState = 'active' | 'aborted' | 'settled';
+    let state: ExecutionState = 'active';
+    let scheduledTimer: ReturnType<typeof setTimeout> | undefined;
+    let scheduledAnimationFrame: number | undefined;
+    let unregisterDeviceLoss = () => false;
+
+    const now = () =>
+      typeof performance === 'undefined' ? Date.now() : performance.now();
+    const executionStartTime = now();
+    const tileTimesMs: number[] = [];
+    const cancelScheduledTile = () => {
+      if (scheduledTimer !== undefined) {
+        clearTimeout(scheduledTimer);
+        scheduledTimer = undefined;
+      }
+      if (
+        scheduledAnimationFrame !== undefined &&
+        typeof cancelAnimationFrame !== 'undefined'
+      ) {
+        cancelAnimationFrame(scheduledAnimationFrame);
+        scheduledAnimationFrame = undefined;
+      }
+    };
+    const reportCallbackFailure = (reason: unknown) => {
+      // 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: unknown) => {
+      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: () => void) => {
+      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);
+      }
+    };
 
-    const executeTile = (i: number, j: number) => {
-      if (aborted) {
+    const executeTile = async (tileIndex: number) => {
+      if (state !== 'active') {
         return;
       }
-      let resGPUBuffer;
-      // profileAndLogKernelCode(() => {
-      ENGINE.startScope();
-      resGPUBuffer = this._executeTile(
+      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,
@@ -750,54 +724,106 @@ class UNet {
           : rawData,
         outputTileData,
         outputImageData,
-        i,
-        j,
+        tile,
+        tileIndex === 0,
         width,
         height,
         hdr,
         denoiseAlpha
       );
-      ENGINE.endScope();
-      // }, true);
+      if (state !== 'active') return;
       const output = outputImageData || {
         data: resGPUBuffer,
         width,
         height
       };
-      progress?.(
-        output as any,
-        // Is undefined if using webgpu buffer
-        outputTileData as any,
-        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);
-          }
-        });
-      } else {
-        // console.log(memory());
-        done(output as any);
+      if (progress) {
+        await progress(
+          output as any,
+          // Is undefined if using webgpu buffer
+          outputTileData as any,
+          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);
+          }
+          // 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 as any);
+          if (state === 'active') {
+            state = 'settled';
+          }
+        }
+      };
+      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();
   }
 }