oidn-web 0.3.5 → 0.5.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 { Tile, type HDRTransfer } from './process';
+import { type DynamicTileSetting } from './tileScheduler';
+import { type UNetModelSpec } from './modelSpec';
+import { type NativeUNetGemmOptions, type NativeUNetKernelSetting, type NativeUNetPrecisionSetting } from './nativeUNet';
+export type UNetEngineSetting = 'auto' | 'wgsl' | 'webnn';
 interface HDRImageData {
     data: Float32Array;
     width: number;
@@ -16,22 +19,41 @@ interface GPUImageDataOutput {
     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;
+    };
+}
 declare class UNet {
-    private _hostTensors;
-    private _backend;
-    private _tfModel;
     private _device;
-    private _tileWidth;
-    private _tileHeight;
-    private _tileOverlapX;
-    private _tileOverlapY;
     private _aux;
     private _hdr;
+    private _hdrTransfer;
     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?;
+    private _activeExecutionFailures;
+    private _deviceLostObserved;
+    private _deviceLostSettled;
+    private _deviceLostReason;
+    constructor(hostTensors: Map<string, HostTensor>, device: GPUDevice, opts?: {
         /**
          * If use auxiliary data.
          */
@@ -40,23 +62,67 @@ declare 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;
     });
-    getDevice(): GPUDevice | undefined;
-    private _buildModel;
-    private _createConv;
-    private _createConcatConv;
-    private _createPooling;
-    private _addUpsamplingLayer;
-    private _addNet;
-    private _addNetLarge;
-    private _updateModel;
-    private _getTileSizeWithOverlap;
+    getDevice(): GPUDevice;
+    /** Completes backend compilation before first interactive use. */
+    prepare(): Promise<void>;
+    /**
+     * Prepares the input shapes selected for an image before its first denoise.
+     * Hosts can call this while they still display their model-loading state.
+     */
+    prepareForImage(width: number, height: number, options?: {
+        tileOverlap?: number;
+        wholeImage?: boolean;
+    }): Promise<void>;
+    getRuntimeInfo(): {
+        configuredEngine: UNetEngineSetting;
+        gpuEngine: "wgsl" | "webnn";
+        precision: import("./nativeUNet").NativeUNetPrecision;
+        kernel: {
+            configured: NativeUNetKernelSetting;
+            gemm: Readonly<Required<NativeUNetGemmOptions>>;
+            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;
+        hdrTransfer: HDRTransfer;
+        dynamicTile: {
+            enabled: boolean;
+            currentTileSize: number;
+            minTileSize: number;
+            maxTileSize: number;
+            targetTileTimeMs: number;
+        };
+        lastExecution: UNetExecutionStats | undefined;
+        activeExecutionCount: number;
+    };
+    private _observeDeviceLoss;
+    private _registerExecutionFailure;
+    /** Captures per-node GPU timestamps for the next native tile execution. */
+    profileNextExecution(): boolean;
+    getLastExecutionProfile(): Promise<import("./nativeUNet").NativeUNetExecutionProfile> | undefined;
     private _processImageData;
     private _readTile;
     private _writeTile;
     private _executeTile;
-    tileExecute<T extends ImageData | HDRImageData | GPUImageData>({ color, albedo, normal, done, progress, denoiseAlpha }: {
+    tileExecute<T extends ImageData | HDRImageData | GPUImageData>({ color, albedo, normal, done, progress, denoiseAlpha, tileOverlap, wholeImage, scheduling, error }: {
         color: T;
         albedo?: ImageData | GPUImageData;
         normal?: ImageData | GPUImageData;
@@ -64,8 +130,27 @@ declare class UNet {
          * If denoise alpha channel. Otherwise denoise RGB channels.
          */
         denoiseAlpha?: boolean;
-        done: (outputData: T extends GPUImageData ? GPUImageDataOutput : T) => void;
-        progress?: (outputData: T extends GPUImageData ? GPUImageDataOutput : T, tileData: (T extends GPUImageData ? GPUImageDataOutput : T) | undefined, tile: Tile, currentIdx: number, totalIdx: number) => 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 | Promise<void>;
     }): () => void;
     dispose(): void;
 }