oidn-webgpu 0.0.0-stage → 0.6.0

package/src/UNet.ts

@@ -0,0 +1,861 @@
+import { HostTensor } from './tza';
+import {
+  GPUDataProcess,
+  Tile,
+  avgLogLum,
+  hdrTransferFuncCPU,
+  hdrTransferFuncInverseCPU,
+  type HDRTransfer
+} from './process';
+import {
+  DynamicTileController,
+  type DynamicTileSetting,
+  planTileGrid,
+  type PlannedTile,
+  OIDN_TILE_ALIGNMENT
+} from './tileScheduler';
+import {
+  detectUNetModelSpec,
+  validateUNetModel,
+  type UNetModelSpec
+} from './modelSpec';
+import {
+  NativeUNetExecutor,
+  type NativeSubgroupMatrixOptions,
+  type NativeUNetGemmOptions,
+  type NativeUNetKernelSetting,
+  type NativeUNetExecutionProfile,
+  type NativeUNetPrecisionSetting
+} from './nativeUNet';
+import { WebNNUNetExecutor } from './webnnUNet';
+
+export type UNetEngineSetting = 'auto' | 'wgsl' | 'webnn';
+
+/** Upper bound on waiting for a display frame between tiles. */
+const ANIMATION_FRAME_FALLBACK_MS = 100;
+
+interface HDRImageData {
+  data: Float32Array;
+  width: number;
+  height: number;
+}
+
+interface GPUImageData {
+  data: GPUBuffer | GPUTexture;
+  width: number;
+  height: number;
+}
+
+interface GPUImageDataOutput {
+  data: GPUBuffer;
+  width: number;
+  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;
+}
+
+function isGPUImageData(
+  data: ImageData | GPUImageData | HDRImageData
+): data is GPUImageData {
+  return data.data instanceof GPUBuffer || data.data instanceof GPUTexture;
+}
+
+class UNet {
+  private _device: GPUDevice;
+
+  private _aux;
+  private _hdr;
+  private _hdrTransfer: HDRTransfer;
+
+  private _dataProcessGPU?: GPUDataProcess;
+  private _nativeExecutor?: NativeUNetExecutor;
+  private _webNNExecutor?: WebNNUNetExecutor;
+  private _modelSpec: UNetModelSpec;
+  private _inputChannels: number;
+  private _engine: UNetEngineSetting;
+
+  private _dynamicTileController: DynamicTileController;
+  private _lastExecution?: UNetExecutionStats;
+  private _profileNextImage = false;
+  private _lastImageProfile?: NativeUNetExecutionProfile;
+  private _activeExecutionFailures = new Set<(reason: unknown) => void>();
+  private _deviceLostObserved = false;
+  private _deviceLostSettled = false;
+  private _deviceLostReason: unknown;
+
+  constructor(
+    hostTensors: Map<string, HostTensor>,
+    device: GPUDevice,
+    opts: {
+      /**
+       * If use auxiliary data.
+       */
+      aux?: boolean;
+      /**
+       * 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;
+      subgroupMatrix?: NativeSubgroupMatrixOptions;
+      /** 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._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,
+          subgroupMatrix: opts.subgroupMatrix
+        }
+      );
+    }
+  }
+
+  getDevice() {
+    return this._device;
+  }
+
+  /** 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;
+    }
+    await this._nativeExecutor!.prepare();
+  }
+
+  /**
+   * Prepares the input shapes selected for an image before its first denoise.
+   * Hosts can call this while they still display their model-loading state.
+   */
+  async prepareForImage(
+    width: number,
+    height: number,
+    options: { tileOverlap?: number; wholeImage?: boolean } = {}
+  ) {
+    const defaultTileOverlap = roundUp(
+      this._modelSpec.receptiveField / 2,
+      OIDN_TILE_ALIGNMENT
+    );
+    const resolvedTileOverlap = options.tileOverlap === undefined
+      ? defaultTileOverlap
+      : roundUp(Math.max(0, options.tileOverlap), OIDN_TILE_ALIGNMENT);
+    const plan = planTileGrid(
+      width,
+      height,
+      options.wholeImage
+        ? roundUp(Math.max(width, height), OIDN_TILE_ALIGNMENT)
+        : this._dynamicTileController.tileSize,
+      resolvedTileOverlap
+    );
+    const shapes = [...new Map(
+      plan.tiles.map(({ input }) => [
+        `${input.width}x${input.height}`,
+        { width: input.width, height: input.height }
+      ])
+    ).values()];
+    await this._webNNExecutor?.prewarm(shapes);
+    this._nativeExecutor?.prewarm(shapes);
+  }
+
+  getRuntimeInfo() {
+    return {
+      configuredEngine: this._engine,
+      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,
+            subgroupMatrix: this._nativeExecutor.subgroupMatrix
+          }
+        : 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 _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 _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);
+  }
+
+  /** Captures network GPU timestamps across every tile of the next image. */
+  profileNextExecution() {
+    if (!this._nativeExecutor || !this._device.features.has('timestamp-query')) return false;
+    this._profileNextImage = true;
+    return true;
+  }
+
+  getLastExecutionProfile() {
+    return this._lastImageProfile ? Promise.resolve(this._lastImageProfile) : undefined;
+  }
+
+  private _processImageData(
+    color: ImageData | HDRImageData,
+    albedo: ImageData | HDRImageData | undefined,
+    normal: ImageData | HDRImageData | undefined,
+    isHDR: boolean
+  ) {
+    const rawData = color.data;
+    const pixelsCount = rawData.length / 4;
+    const channels = this._aux ? 9 : 3;
+    const tensorData = new Float32Array(pixelsCount * channels);
+
+    if ((albedo && !normal) || (normal && !albedo)) {
+      throw new Error('Normal map and albedo map are both required');
+    }
+    if (albedo && normal) {
+      if (
+        albedo.width !== normal.width ||
+        albedo.height !== normal.height ||
+        color.width !== albedo.width ||
+        color.height !== albedo.height
+      ) {
+        throw new Error('Image size mismatch');
+      }
+    }
+
+    const albedoData = albedo?.data;
+    const normalData = normal?.data;
+    // Each guide can be an 8-bit image or normalized float RGBA,
+    // independently of the color image's HDR setting and the other guide.
+    const albedoDivisor = albedoData instanceof Float32Array ? 1 : 255;
+    const normalDivisor = normalData instanceof Float32Array ? 1 : 255;
+    for (let i = 0; i < rawData.length; i += 4) {
+      const i2 = (i / 4) * channels;
+
+      for (let c = 0; c < 3; c++) {
+        if (isHDR) {
+          tensorData[i2 + c] = rawData[i + c];
+        } else {
+          tensorData[i2 + c] = rawData[i + c] / 255;
+        }
+        if (albedoData) {
+          tensorData[i2 + c + 3] = albedoData[i + c] / albedoDivisor;
+        }
+        if (normalData) {
+          tensorData[i2 + c + 6] = normalData[i + c] / normalDivisor;
+        }
+      }
+    }
+
+    return tensorData;
+  }
+
+  private _readTile(
+    data: Float32Array,
+    channels: number,
+    srcTile: Tile,
+    width: number
+  ) {
+    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 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];
+        }
+      }
+    }
+    return tileData;
+  }
+
+  private _writeTile(
+    imageData: ImageData | HDRImageData,
+    srcTile: Tile,
+    dstTile: Tile,
+    srcTileData: Float32Array,
+    srcWidth: number,
+    isHDR: boolean
+  ) {
+    const { data: outImageData, width } = imageData;
+    const dx = dstTile.x - srcTile.x;
+    const dy = dstTile.y - srcTile.y;
+    for (let y = 0; y < dstTile.height; y++) {
+      for (let x = 0; x < dstTile.width; x++) {
+        const i1 = ((y + dy) * srcWidth + x + dx) * 3;
+        const i2 = ((y + dstTile.y) * width + (x + dstTile.x)) * 4;
+
+        for (let c = 0; c < 3; c++) {
+          if (isHDR) {
+            outImageData[i2 + c] = srcTileData[i1 + c];
+          } else {
+            outImageData[i2 + c] = Math.min(
+              Math.max(srcTileData[i1 + c] * 255, 0),
+              255
+            );
+          }
+        }
+        imageData.data[i2 + 3] = isHDR ? 1 : 255;
+      }
+    }
+  }
+
+  private async _executeTile(
+    inputData:
+      | Float32Array
+      | {
+          color: GPUBuffer | GPUTexture;
+          // TODO optional
+          albedo?: GPUBuffer | GPUTexture;
+          normal?: GPUBuffer | GPUTexture;
+        },
+    outputTileData: ImageData | HDRImageData | undefined,
+    outputImageData: ImageData | HDRImageData | undefined,
+    tile: PlannedTile,
+    isFirstTile: boolean,
+    width: number,
+    height: number,
+    isHDR: boolean,
+    denoiseAlpha?: boolean
+  ) {
+    const channels = this._aux ? 9 : 3;
+    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: GPUBuffer | undefined;
+    let denoisedData: Float32Array | undefined;
+    let inputScale = 1;
+    const device = this._device;
+    let dataProcessGPU = this._dataProcessGPU;
+
+    if (inputData instanceof Float32Array) {
+      let tileData = this._readTile(inputData, channels, srcTile, width);
+      if (isHDR) {
+        inputScale = avgLogLum({
+          data: tileData,
+          channels
+        });
+        tileData = hdrTransferFuncCPU({
+          data: tileData,
+          channels,
+          inputScale,
+          transfer: this._hdrTransfer
+        });
+      }
+      denoisedData = await (this._webNNExecutor ?? this._nativeExecutor!).executeCPU(
+        tileData,
+        srcTileWidth,
+        srcTileHeight
+      );
+    } else {
+      if (!dataProcessGPU) {
+        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 (isFirstTile) {
+        dataProcessGPU.copyInputDataToOutput(inputData.color);
+      }
+      const { color, albedo, normal } = dataProcessGPU.forward(
+        inputData.color,
+        this._aux ? inputData.albedo : undefined,
+        this._aux ? inputData.normal : undefined,
+        denoiseAlpha
+      );
+
+      nativeOutputBuffer = await (this._webNNExecutor ?? this._nativeExecutor!).execute(
+        this._aux ? [color, albedo!, normal!] : [color],
+        srcTileWidth,
+        srcTileHeight
+      );
+    }
+
+    let outBuffer: GPUBuffer;
+
+    if (inputData instanceof Float32Array) {
+      if (isHDR) {
+        denoisedData = hdrTransferFuncInverseCPU({
+          data: denoisedData!,
+          channels: 3,
+          inputScale,
+          transfer: this._hdrTransfer
+        });
+      }
+
+      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];
+          }
+        }
+      }
+    } else {
+      dataProcessGPU!.setOutputTile(dstTile, srcTile);
+      outBuffer = dataProcessGPU!.inverse(
+        nativeOutputBuffer!,
+        inputData.color
+      );
+    }
+    return outBuffer!;
+  }
+
+  tileExecute<T extends ImageData | HDRImageData | GPUImageData>({
+    color,
+    albedo,
+    normal,
+    done,
+    progress,
+    denoiseAlpha,
+    tileOverlap,
+    wholeImage,
+    scheduling = 'event-loop',
+    error
+  }: {
+    color: T;
+    albedo?: ImageData | HDRImageData | GPUImageData;
+    normal?: ImageData | HDRImageData | GPUImageData;
+    /**
+     * If denoise alpha channel. Otherwise denoise RGB channels.
+     */
+    denoiseAlpha?: boolean;
+    /**
+     * Execute the complete input image as one tile, ignoring `maxTileSize`.
+     * The image must fit the device's buffer and dispatch limits.
+     */
+    wholeImage?: boolean;
+    /**
+     * Per-side context for boundaries shared with another tile. Defaults to
+     * half of the model receptive field rounded up to 16 pixels.
+     */
+    tileOverlap?: number;
+    /**
+     * How JavaScript yields between completed GPU tiles. `event-loop`
+     * (default) continues on the next macrotask. `animation-frame` waits for
+     * the next display frame, bounded by a short timer so hidden pages still
+     * complete.
+     */
+    scheduling?: 'animation-frame' | 'event-loop';
+    done: (
+      outputData: T extends GPUImageData ? GPUImageDataOutput : T
+    ) => void | Promise<void>;
+    /** Receives asynchronous execution, queue, and callback failures. */
+    error?: (reason: unknown) => void | Promise<void>;
+    progress?: (
+      outputData: T extends GPUImageData ? GPUImageDataOutput : T,
+      tileData: (T extends GPUImageData ? GPUImageDataOutput : T) | undefined,
+      tile: Tile,
+      currentIdx: number,
+      totalIdx: number
+    ) => void | Promise<void>;
+  }): () => void {
+    if (this._aux && (!albedo || !normal)) {
+      throw new Error('Normal map and albedo map are both required');
+    }
+
+    if (!this._aux) {
+      if (albedo || normal) {
+        throw new Error('Normal map and albedo map are not required');
+      }
+    }
+
+    const width = color.width;
+    const height = color.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;
+    let rawData: Float32Array;
+    if (!isGPUImageData(color)) {
+      rawData = this._processImageData(
+        color,
+        albedo as ImageData | HDRImageData | undefined,
+        normal as ImageData | HDRImageData | undefined,
+        hdr
+      );
+    }
+    function makeImageData(width: number, height: number) {
+      return hdr
+        ? {
+            data: new Float32Array(width * height * 4),
+            width,
+            height
+          }
+        : new ImageData(width, height);
+    }
+
+    const outputImageData = isGPUImageData(color)
+      ? undefined
+      : makeImageData(width, height);
+
+    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 profileImage = this._profileNextImage === true;
+    this._profileNextImage = false;
+    const layerTimes = new Map<string, 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 = async (tileIndex: number) => {
+      if (state !== 'active') {
+        return;
+      }
+      const tile = plan.tiles[tileIndex];
+      const outputTileData = isGPUImageData(color)
+        ? undefined
+        : makeImageData(tile.output.width, tile.output.height);
+      const tileStartTime = now();
+      if (profileImage) this._nativeExecutor!.profileNextExecution();
+      const resGPUBuffer = await this._executeTile(
+        isGPUImageData(color)
+          ? {
+              color: color.data,
+              albedo: (albedo as GPUImageData | undefined)?.data,
+              normal: (normal as GPUImageData | undefined)?.data
+            }
+          : rawData,
+        outputTileData,
+        outputImageData,
+        tile,
+        tileIndex === 0,
+        width,
+        height,
+        hdr,
+        denoiseAlpha
+      );
+      if (profileImage) {
+        const profile = await this._nativeExecutor!.getLastExecutionProfile();
+        for (const layer of profile?.layers ?? []) {
+          layerTimes.set(layer.id, (layerTimes.get(layer.id) ?? 0) + layer.durationMs);
+        }
+      }
+      if (state !== 'active') return;
+      const output = outputImageData || {
+        data: resGPUBuffer,
+        width,
+        height
+      };
+      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();
+          if (profileImage) {
+            const layers = [...layerTimes].map(([id, durationMs]) => ({ id, durationMs }));
+            this._lastImageProfile = { totalMs: layers.reduce((sum, layer) => sum + layer.durationMs, 0), layers };
+          }
+          await done(output as any);
+          if (state === 'active') {
+            state = 'settled';
+          }
+        }
+      };
+      await continueAfterGPUWork();
+    };
+
+    unregisterDeviceLoss = this._registerExecutionFailure(settleError);
+    void executeTile(0).catch(settleError);
+
+    return () => {
+      if (state !== 'active') return;
+      state = 'aborted';
+      cancelScheduledTile();
+      unregisterDeviceLoss();
+    };
+  }
+
+  dispose() {
+    this._dataProcessGPU?.dispose();
+    this._nativeExecutor?.dispose();
+    this._webNNExecutor?.dispose();
+  }
+}
+
+export default UNet;