oidn-web 0.4.0 → 0.5.0

package/src/UNet.ts

@@ -4,14 +4,15 @@ import {
   Tile,
   avgLogLum,
   hdrTransferFuncCPU,
-  hdrTransferFuncInverseCPU
+  hdrTransferFuncInverseCPU,
+  type HDRTransfer
 } from './process';
 import {
   DynamicTileController,
   type DynamicTileSetting,
-  fitTileDimension,
-  OIDN_TILE_ALIGNMENT,
-  waitForSubmittedGPUWork
+  planTileGrid,
+  type PlannedTile,
+  OIDN_TILE_ALIGNMENT
 } from './tileScheduler';
 import {
   detectUNetModelSpec,
@@ -20,6 +21,7 @@ import {
 } from './modelSpec';
 import {
   NativeUNetExecutor,
+  type NativeUNetGemmOptions,
   type NativeUNetKernelSetting,
   type NativeUNetPrecisionSetting
 } from './nativeUNet';
@@ -27,6 +29,9 @@ 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;
@@ -48,9 +53,12 @@ interface GPUImageDataOutput {
 export interface UNetExecutionStats {
   width: number;
   height: number;
-  tileWidth: number;
-  tileHeight: number;
   tileCount: number;
+  tileColumns: number;
+  tileRows: number;
+  tileOverlap: number;
+  inputPixelCount: number;
+  inputShapeCount: number;
   durationMs: number;
   tileTimeMs: {
     min: number;
@@ -71,18 +79,11 @@ function isGPUImageData(
 }
 
 class UNet {
-  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 _nativeExecutor?: NativeUNetExecutor;
@@ -93,10 +94,14 @@ class UNet {
 
   private _dynamicTileController: DynamicTileController;
   private _lastExecution?: UNetExecutionStats;
+  private _activeExecutionFailures = new Set<(reason: unknown) => void>();
+  private _deviceLostObserved = false;
+  private _deviceLostSettled = false;
+  private _deviceLostReason: unknown;
 
   constructor(
     hostTensors: Map<string, HostTensor>,
-    backend: { device: GPUDevice; adapterInfo: GPUAdapterInfo },
+    device: GPUDevice,
     opts: {
       /**
        * If use auxiliary data.
@@ -106,20 +111,24 @@ class UNet {
        * If input is HDR image.
        */
       hdr?: boolean;
+      /** HDR transfer function expected by the trained model. */
+      hdrTransfer?: HDRTransfer;
       maxTileSize?: number;
       dynamicTile?: DynamicTileSetting;
-      /** Reserved for explicit native WGSL selection. */
+      /** Native WGSL or the experimental WebNN backend. */
       engine?: UNetEngineSetting;
-      /** Arithmetic/storage precision used by the native WGSL engine. */
+      /** Arithmetic/storage precision used by the native WGSL executor. */
       precision?: NativeUNetPrecisionSetting;
       /** Model-independent convolution kernel selection. */
       kernel?: NativeUNetKernelSetting;
+      gemm?: NativeUNetGemmOptions;
       /** Explicit descriptor for a new OIDN topology not in the built-in registry. */
       modelSpec?: UNetModelSpec;
     } = {}
   ) {
     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);
@@ -139,7 +148,8 @@ class UNet {
       opts.dynamicTile
     );
 
-    this._device = backend.device;
+    this._device = device;
+    this._observeDeviceLoss();
     if (this._engine === 'webnn') {
       this._webNNExecutor = new WebNNUNetExecutor(
         this._device,
@@ -150,7 +160,7 @@ class UNet {
       this._nativeExecutor = new NativeUNetExecutor(
         this._device,
         validatedModel,
-        { precision: opts.precision, kernel: opts.kernel }
+        { precision: opts.precision, kernel: opts.kernel, gemm: opts.gemm }
       );
     }
   }
@@ -182,6 +192,40 @@ class UNet {
     await this._nativeExecutor!.prepare();
   }
 
+  /**
+   * Prepares the input shapes selected for an image before its first denoise.
+   * Hosts can call this while they still display their model-loading state.
+   */
+  async prepareForImage(
+    width: 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,
@@ -190,6 +234,7 @@ class UNet {
       kernel: this._nativeExecutor
         ? {
             configured: this._nativeExecutor.kernelSetting,
+            gemm: this._nativeExecutor.gemm,
             maxSpatialInputBlocks: this._nativeExecutor.maxSpatialInputBlocks,
             subgroupsAvailable: this._nativeExecutor.subgroupsAvailable
           }
@@ -201,6 +246,7 @@ class UNet {
       model: this._modelSpec.id,
       modelFamily: this._modelSpec.family,
       inputChannels: this._inputChannels,
+      hdrTransfer: this._hdrTransfer,
       dynamicTile: {
         enabled: this._dynamicTileController.enabled,
         currentTileSize: this._dynamicTileController.tileSize,
@@ -208,10 +254,44 @@ class UNet {
         maxTileSize: this._dynamicTileController.maxTileSize,
         targetTileTimeMs: this._dynamicTileController.targetTileTimeMs
       },
-      lastExecution: this._lastExecution
+      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 per-node GPU timestamps for the next native tile execution. */
   profileNextExecution() {
     return this._nativeExecutor?.profileNextExecution() ?? false;
@@ -221,54 +301,6 @@ class UNet {
     return this._nativeExecutor?.getLastExecutionProfile();
   }
 
-  private _updateModel(width: number, height: number) {
-    const maxTileSize = this._dynamicTileController.tileSize;
-
-    let tileWidth = fitTileDimension(width, maxTileSize);
-    let tileHeight = fitTileDimension(height, maxTileSize);
-    const defaultTileOverlap = roundUp(
-      this._modelSpec.receptiveField / 2,
-      OIDN_TILE_ALIGNMENT
-    );
-    let tileOverlapX = defaultTileOverlap;
-    let tileOverlapY = defaultTileOverlap;
-
-    if (width <= maxTileSize) {
-      tileOverlapX = 0;
-    }
-    if (height <= maxTileSize) {
-      tileOverlapY = 0;
-    }
-
-    // Force width and height has same size. reduce the cache in memory
-    const tileSize = Math.max(tileWidth, tileHeight);
-    const tileOverlap = Math.max(tileOverlapX, tileOverlapY);
-    tileWidth = tileSize;
-    tileHeight = tileSize;
-    tileOverlapX = tileOverlap;
-    tileOverlapY = tileOverlap;
-
-    if (
-      tileWidth !== this._tileWidth ||
-      tileHeight !== this._tileHeight ||
-      tileOverlapX !== this._tileOverlapX ||
-      tileOverlapY !== this._tileOverlapY
-    ) {
-      // console.log(tileWidth, tileHeight, tileOverlapX, tileOverlapY);
-      this._tileWidth = tileWidth;
-      this._tileHeight = tileHeight;
-      this._tileOverlapX = tileOverlapX;
-      this._tileOverlapY = tileOverlapY;
-    }
-  }
-
-  private _getTileSizeWithOverlap() {
-    return {
-      width: this._tileWidth + 2 * this._tileOverlapX,
-      height: this._tileHeight + 2 * this._tileOverlapY
-    };
-  }
-
   private _processImageData(
     color: ImageData | HDRImageData,
     albedo: ImageData | undefined,
@@ -326,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++) {
@@ -381,36 +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 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) {
@@ -423,7 +455,8 @@ class UNet {
         tileData = hdrTransferFuncCPU({
           data: tileData,
           channels,
-          inputScale
+          inputScale,
+          transfer: this._hdrTransfer
         });
       }
       denoisedData = await (this._webNNExecutor ?? this._nativeExecutor!).executeCPU(
@@ -435,13 +468,14 @@ class UNet {
       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(
@@ -460,18 +494,13 @@ class UNet {
 
     let outBuffer: GPUBuffer;
 
-    const dstWidth = Math.min(dstTileSize.width, width);
-    const dstHeight = Math.min(dstTileSize.height, height);
-    const dstTile = new Tile(i * dstWidth, j * dstHeight, dstWidth, dstHeight);
-    dstTile.width = Math.min(dstTile.width, width - dstTile.x);
-    dstTile.height = Math.min(dstTile.height, height - dstTile.y);
-
     if (inputData instanceof Float32Array) {
       if (isHDR) {
         denoisedData = hdrTransferFuncInverseCPU({
           data: denoisedData!,
           channels: 3,
-          inputScale
+          inputScale,
+          transfer: this._hdrTransfer
         });
       }
 
@@ -480,13 +509,13 @@ class UNet {
         srcTile,
         dstTile,
         denoisedData!,
-        srcTileSize.width,
+        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];
@@ -509,7 +538,11 @@ class UNet {
     normal,
     done,
     progress,
-    denoiseAlpha
+    denoiseAlpha,
+    tileOverlap,
+    wholeImage,
+    scheduling = 'event-loop',
+    error
   }: {
     color: T;
     albedo?: ImageData | GPUImageData;
@@ -518,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');
@@ -540,9 +594,24 @@ class UNet {
     const width = color.width;
     const height = color.height;
     const adaptiveTileSize = this._dynamicTileController.tileSize;
-    const shouldAdaptTileSize =
-      width > adaptiveTileSize || height > adaptiveTileSize;
-    this._updateModel(width, height);
+    // The planner aligns the maximum down, so round up to keep one tile.
+    const requestedTileSize = wholeImage
+      ? roundUp(Math.max(width, height), OIDN_TILE_ALIGNMENT)
+      : adaptiveTileSize;
+    const defaultTileOverlap = roundUp(
+      this._modelSpec.receptiveField / 2,
+      OIDN_TILE_ALIGNMENT
+    );
+    const resolvedTileOverlap = tileOverlap === undefined
+      ? defaultTileOverlap
+      : roundUp(Math.max(0, tileOverlap), OIDN_TILE_ALIGNMENT);
+    const plan = planTileGrid(
+      width,
+      height,
+      requestedTileSize,
+      resolvedTileOverlap
+    );
+    const shouldAdaptTileSize = plan.tiles.length > 1;
 
     // TODO should fixed to be hdr when UNet is created.
     // weights of hdr and ldr is different
@@ -557,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
         ? {
@@ -575,28 +639,80 @@ 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 (typeof requestAnimationFrame === 'undefined') {
-        setTimeout(callback, 0);
+      if (
+        scheduling === 'event-loop' ||
+        typeof requestAnimationFrame === 'undefined'
+      ) {
+        scheduledTimer = setTimeout(() => {
+          scheduledTimer = undefined;
+          callback();
+        }, 0);
       } else {
-        requestAnimationFrame(callback);
+        // Hidden documents pause requestAnimationFrame. Race it against a
+        // timer so animation-frame scheduling still completes in background
+        // tabs, minimized windows, and offscreen iframes.
+        const run = () => {
+          cancelScheduledTile();
+          callback();
+        };
+        scheduledAnimationFrame = requestAnimationFrame(run);
+        scheduledTimer = setTimeout(run, ANIMATION_FRAME_FALLBACK_MS);
       }
     };
 
-    const executeTile = async (i: number, j: number) => {
-      if (aborted) {
+    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();
       const resGPUBuffer = await this._executeTile(
         isGPUImageData(color)
@@ -608,41 +724,47 @@ class UNet {
           : rawData,
         outputTileData,
         outputImageData,
-        i,
-        j,
+        tile,
+        tileIndex === 0,
         width,
         height,
         hdr,
         denoiseAlpha
       );
-      if (aborted) return;
+      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 (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 = i + 1 < tileCountW || j + 1 < tileCountH;
-      const continueAfterGPUWork = () => {
+      const hasNextTile = tileIndex + 1 < plan.tiles.length;
+      await this._device.queue.onSubmittedWorkDone();
+      if (state !== 'active') return;
+      const continueAfterGPUWork = async () => {
         tileTimesMs.push(now() - tileStartTime);
-        if (aborted) return;
+        if (state !== 'active') return;
 
         if (hasNextTile) {
           scheduleNextTile(() => {
-            if (aborted) return;
-            if (i + 1 < tileCountW) {
-              executeTile(i + 1, j);
-            } else if (j + 1 < tileCountH) {
-              executeTile(0, j + 1);
-            }
+            if (state !== 'active') return;
+            void executeTile(tileIndex + 1).catch(settleError);
           });
         } else {
           const sortedTileTimes = [...tileTimesMs].sort((a, b) => a - b);
@@ -653,9 +775,12 @@ class UNet {
           this._lastExecution = {
             width,
             height,
-            tileWidth,
-            tileHeight,
-            tileCount: tileCountW * tileCountH,
+            tileCount: plan.tiles.length,
+            tileColumns: plan.columns,
+            tileRows: plan.rows,
+            tileOverlap: plan.overlap,
+            inputPixelCount: plan.inputPixelCount,
+            inputShapeCount: plan.inputShapeCount,
             durationMs: now() - executionStartTime,
             tileTimeMs: {
               min: sortedTileTimes[0],
@@ -671,23 +796,27 @@ class UNet {
           if (shouldAdaptTileSize) {
             this._dynamicTileController.observe(tileTimesMs);
           }
-          // console.log(memory());
-          done(output as any);
+          // 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';
+          }
         }
       };
-
-      // requestAnimationFrame only throttles JavaScript submission. Waiting
-      // for the queue here keeps at most one OIDN tile in flight, so aborting
-      // cannot leave a long tail of already-submitted GPU work.
-      void waitForSubmittedGPUWork(this._device!.queue).then(
-        continueAfterGPUWork
-      );
+      await continueAfterGPUWork();
     };
 
-    executeTile(0, 0);
+    unregisterDeviceLoss = this._registerExecutionFailure(settleError);
+    void executeTile(0).catch(settleError);
 
     return () => {
-      aborted = true;
+      if (state !== 'active') return;
+      state = 'aborted';
+      cancelScheduledTile();
+      unregisterDeviceLoss();
     };
   }