oidn-web 0.2.2 → 0.3.0

package/src/UNet.ts

@@ -7,6 +7,7 @@ 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
@@ -25,7 +26,8 @@ import {
   hdrTransferFuncInverseCPU
 } from './process';
 import type { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu';
-// import { profileAndLogKernelCode } from './helper';
+
+// import { profileAndLogKernelCode, memory } from './helper';
 
 function getTensorData(
   ubytes: Uint8Array,
@@ -68,6 +70,12 @@ interface HDRImageData {
 }
 
 interface GPUImageData {
+  data: GPUBuffer | GPUTexture;
+  width: number;
+  height: number;
+}
+
+interface GPUImageDataOutput {
   data: GPUBuffer;
   width: number;
   height: number;
@@ -84,7 +92,7 @@ function roundUp2(a: number, b: number, c: number) {
 function isGPUImageData(
   data: ImageData | GPUImageData | HDRImageData
 ): data is GPUImageData {
-  return data.data instanceof GPUBuffer;
+  return data.data instanceof GPUBuffer || data.data instanceof GPUTexture;
 }
 
 const receptiveField = 174; // receptive field in pixels
@@ -109,18 +117,27 @@ class UNet {
   private _tileOverlapX = 0;
   private _tileOverlapY = 0;
 
-  private _aux = false;
-  private _hdr = false;
+  private _aux;
+  private _hdr;
 
   private _dataProcessGPU?: GPUDataProcess;
 
   private _maxTileSize;
 
+  private _tensors = new Map<string, Tensor>();
+  private _modelsCache = new Map<string, LayersModel>();
+
   constructor(
-    private _tensors: Map<string, HostTensor>,
+    private _hostTensors: Map<string, HostTensor>,
     private _backend: WebGPUBackend,
     opts: {
+      /**
+       * If use auxiliary data.
+       */
       aux?: boolean;
+      /**
+       * If input is HDR image.
+       */
       hdr?: boolean;
       maxTileSize?: number;
     } = {}
@@ -128,7 +145,7 @@ class UNet {
     this._aux = opts.aux || false;
     this._hdr = opts.hdr || false;
 
-    this._maxTileSize = opts.maxTileSize ?? 512;
+    this._maxTileSize = roundUp(opts.maxTileSize ?? 512, 2);
 
     this._device = this._backend.device;
   }
@@ -141,19 +158,30 @@ class UNet {
     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;
+    }
 
-    // TODO input process transferFunc
-    // TODO input shape
     const input = TFInput({
+      name: 'input',
       shape: [tileSize.height, tileSize.width, channels],
       dtype: 'float32'
     });
 
     this._tfModel = new LayersModel({
       inputs: [input],
-      // TODO output process transferFunc
       outputs: isLarge ? this._addNetLarge(input) : this._addNet(input)
     });
+    cache.set(key, this._tfModel);
   }
 
   private _createConv(
@@ -161,21 +189,34 @@ class UNet {
     source: SymbolicTensor,
     activation?: 'relu'
   ) {
-    const unetWeightTensor = this._tensors.get(name + '.weight')!;
-    const unetBiasTensor = this._tensors.get(name + '.bias')!;
-    const weightDims = unetWeightTensor.desc.dims;
-    const weightTensor = tensor(
-      changeWeightShapes(
-        getTensorData(unetWeightTensor.data, unetWeightTensor.desc.dataType),
-        weightDims
-      ),
-      [weightDims[2], weightDims[3], weightDims[1], weightDims[0]],
-      'float32'
-    );
-    const biasTensor = tensor1d(
-      getTensorData(unetBiasTensor.data, unetBiasTensor.desc.dataType),
-      'float32'
-    );
+    const weightTensorName = name + '.weight';
+    const biasTensorName = name + '.bias';
+    const tensors = this._tensors;
+    let weightTensor = tensors.get(weightTensorName);
+    let biasTensor = tensors.get(biasTensorName);
+    const unetWeightTensor = this._hostTensors.get(weightTensorName)!;
+
+    if (!weightTensor) {
+      const weightDims = unetWeightTensor.desc.dims;
+      weightTensor = tensor(
+        changeWeightShapes(
+          getTensorData(unetWeightTensor.data, unetWeightTensor.desc.dataType),
+          weightDims
+        ),
+        [weightDims[2], weightDims[3], weightDims[1], weightDims[0]],
+        'float32'
+      );
+
+      tensors.set(weightTensorName, weightTensor);
+    }
+    if (!biasTensor) {
+      const unetBiasTensor = this._hostTensors.get(name + '.bias')!;
+      biasTensor = tensor1d(
+        getTensorData(unetBiasTensor.data, unetBiasTensor.desc.dataType),
+        'float32'
+      );
+      tensors.set(biasTensorName, biasTensor);
+    }
     // TODO whats the purpose of padded dims ?
     const convLayer = new Conv2D({
       name,
@@ -196,36 +237,39 @@ class UNet {
     source1: SymbolicTensor,
     source2: SymbolicTensor
   ) {
+    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
-      new Concatenate({ trainable: false, axis: 3 }).apply([
-        // convLayer.apply(source2) as SymbolicTensor,
-        source1,
-        source2
-      ]) as SymbolicTensor,
+      concatLayer.apply([source1, source2]) as SymbolicTensor,
       'relu'
     ) as SymbolicTensor;
   }
 
   private _createPooling(source: SymbolicTensor) {
-    // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/training/model.py#L33
-    return new MaxPooling2D({
+    const poolingLayer = new MaxPooling2D({
       name: source.name + '/pooling',
       poolSize: [2, 2],
       strides: [2, 2],
       padding: 'same',
       trainable: false
-    }).apply(source) as SymbolicTensor;
+    });
+    // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/training/model.py#L33
+    return poolingLayer.apply(source) as SymbolicTensor;
   }
 
   private _addUpsamplingLayer(source: SymbolicTensor) {
-    return new UpSampling2D({
+    const upsamplingLayer = new UpSampling2D({
       name: source.name + '/upsampling',
       size: [2, 2],
       trainable: false
-    }).apply(source) as SymbolicTensor;
+    });
+    return upsamplingLayer.apply(source) as SymbolicTensor;
   }
 
   private _addNet(input: SymbolicTensor) {
@@ -308,14 +352,26 @@ class UNet {
     let tileOverlapY = isLarge ? defaultTileOverlapLarge : defaultTileOverlap;
 
     if (width < maxTileSize + defaultTileOverlap * 2) {
-      tileWidth = roundUp(width, tileAlignment);
-      tileOverlapX = 0;
+      tileWidth = roundUp(width, maxTileSize / 2);
+      if (width <= maxTileSize) {
+        tileOverlapX = 0;
+      }
     }
     if (height < maxTileSize + defaultTileOverlap * 2) {
-      tileHeight = roundUp(height, tileAlignment);
-      tileOverlapY = 0;
+      tileHeight = roundUp(height, maxTileSize / 2);
+      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 ||
@@ -323,15 +379,12 @@ class UNet {
       tileOverlapY !== this._tileOverlapY ||
       !this._tfModel
     ) {
+      // console.log(tileWidth, tileHeight, tileOverlapX, tileOverlapY);
       this._tileWidth = tileWidth;
       this._tileHeight = tileHeight;
       this._tileOverlapX = tileOverlapX;
       this._tileOverlapY = tileOverlapY;
 
-      if (this._tfModel) {
-        this._tfModel.dispose();
-      }
-
       this._buildModel(isLarge);
     }
   }
@@ -448,10 +501,10 @@ class UNet {
     inputData:
       | Float32Array
       | {
-          color: GPUBuffer;
+          color: GPUBuffer | GPUTexture;
           // TODO optional
-          albedo: GPUBuffer;
-          normal: GPUBuffer;
+          albedo?: GPUBuffer | GPUTexture;
+          normal?: GPUBuffer | GPUTexture;
         },
     outputTileData: ImageData | HDRImageData | undefined,
     outputImageData: ImageData | HDRImageData | undefined,
@@ -459,7 +512,8 @@ class UNet {
     j: number,
     width: number,
     height: number,
-    isHDR: boolean
+    isHDR: boolean,
+    denoiseAlpha?: boolean
   ) {
     const channels = this._aux ? 9 : 3;
     const tileOverlapX = this._tileOverlapX;
@@ -475,10 +529,9 @@ class UNet {
     let srcY1 = Math.min(srcY0 + srcTileSize.height, height);
     srcY0 = Math.max(srcY1 - srcTileSize.height, 0);
 
-    const srcTileWidth = Math.min(srcTileSize.width, width);
-    const srcTileHeight = Math.min(srcTileSize.height, height);
-    const needsResize =
-      width < dstTileSize.width || height < dstTileSize.height;
+    const srcTileWidth = srcTileSize.width;
+    const srcTileHeight = srcTileSize.height;
+
     const srcTile = new Tile(srcX0, srcY0, srcTileWidth, srcTileHeight);
 
     let tileTensor!: Tensor;
@@ -491,11 +544,11 @@ class UNet {
       if (isHDR) {
         inputScale = avgLogLum({
           data: tileData,
-          channels: 9
+          channels
         });
         tileData = hdrTransferFuncCPU({
           data: tileData,
-          channels: 9,
+          channels,
           inputScale
         });
       }
@@ -505,11 +558,11 @@ class UNet {
         'float32'
       ) as Tensor4D;
     } else {
-      if (!isHDR) {
-        throw new Error('Only hdr is supported for webgpu data.');
-      }
       if (!dataProcessGPU) {
-        dataProcessGPU = this._dataProcessGPU = new GPUDataProcess(device);
+        dataProcessGPU = this._dataProcessGPU = new GPUDataProcess(
+          device,
+          isHDR
+        );
       }
       dataProcessGPU.setImageSize(width, height);
       dataProcessGPU.setInputTile(srcTile);
@@ -519,43 +572,38 @@ class UNet {
       }
       const { color, albedo, normal } = dataProcessGPU.forward(
         inputData.color,
-        inputData.albedo,
-        inputData.normal
+        this._aux ? inputData.albedo : undefined,
+        this._aux ? inputData.normal : undefined,
+        denoiseAlpha
       );
-      const shape = [1, srcTileHeight, srcTileWidth, 4] as any;
-      const auxTensors = [color, albedo, normal].map((buffer) => {
-        const tmp = tensor({ buffer, zeroCopy: true }, shape) as Tensor4D;
+
+      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]
         );
-        tmp.dispose();
         return ret;
-      });
+      };
 
-      tileTensor = concat4d(auxTensors, 3);
-      // TODO reuse tensors?
-      auxTensors.forEach((t) => t.dispose());
-    }
-    // We need resize if input size is smaller than tile size. And is rounded up.
-    if (needsResize) {
-      const rawTileTensor = tileTensor;
-      tileTensor = mirrorPad(
-        rawTileTensor,
-        [
-          [0, 0],
-          [0, srcTileSize.height - height],
-          [0, srcTileSize.width - width],
-          [0, 0]
-        ],
-        'reflect'
-      );
-      rawTileTensor.dispose();
+      if (this._aux) {
+        const tensors = [color, albedo, normal].map((buffer) =>
+          createTensor(buffer!)
+        );
+        tileTensor = concat4d(tensors, 3);
+      } else {
+        tileTensor = createTensor(color);
+      }
     }
 
+    let outBuffer: GPUBuffer;
     const outputTensor = this._tfModel!.predict(tileTensor) as Tensor;
-    tileTensor.dispose();
 
     const dstWidth = Math.min(dstTileSize.width, width);
     const dstHeight = Math.min(dstTileSize.height, height);
@@ -591,8 +639,6 @@ class UNet {
           }
         }
       }
-
-      outputTensor.dispose();
     } else {
       dataProcessGPU!.setOutputTile(dstTile, srcTile);
       // IMPORTANT
@@ -604,14 +650,12 @@ class UNet {
         [0, 0],
         [0, 1]
       ]);
-      const outBuffer = dataProcessGPU!.inverse(
+      outBuffer = dataProcessGPU!.inverse(
         outputTensor4Channnels.dataToGPU().buffer!,
         inputData.color
       );
-      outputTensor.dispose();
-      outputTensor4Channnels.dispose();
-      return outBuffer;
     }
+    return outBuffer!;
   }
 
   tileExecute<T extends ImageData | HDRImageData | GPUImageData>({
@@ -619,15 +663,20 @@ class UNet {
     albedo,
     normal,
     done,
-    progress
+    progress,
+    denoiseAlpha
   }: {
     color: T;
     albedo?: ImageData | GPUImageData;
     normal?: ImageData | GPUImageData;
-    done: (outputData: T) => void;
+    /**
+     * If denoise alpha channel. Otherwise denoise RGB channels.
+     */
+    denoiseAlpha?: boolean;
+    done: (outputData: T extends GPUImageData ? GPUImageDataOutput : T) => void;
     progress?: (
-      outputData: T,
-      tileData: T | undefined,
+      outputData: T extends GPUImageData ? GPUImageDataOutput : T,
+      tileData: (T extends GPUImageData ? GPUImageDataOutput : T) | undefined,
       tile: Tile,
       currentIdx: number,
       totalIdx: number
@@ -655,8 +704,8 @@ class UNet {
     if (!isGPUImageData(color)) {
       rawData = this._processImageData(
         color,
-        albedo as ImageData,
-        normal as ImageData,
+        albedo as ImageData | undefined,
+        normal as ImageData | undefined,
         hdr
       );
     }
@@ -690,12 +739,13 @@ class UNet {
       }
       let resGPUBuffer;
       // profileAndLogKernelCode(() => {
+      ENGINE.startScope();
       resGPUBuffer = this._executeTile(
         isGPUImageData(color)
           ? {
               color: color.data,
-              albedo: (albedo as GPUImageData).data,
-              normal: (normal as GPUImageData).data
+              albedo: (albedo as GPUImageData | undefined)?.data,
+              normal: (normal as GPUImageData | undefined)?.data
             }
           : rawData,
         outputTileData,
@@ -704,8 +754,10 @@ class UNet {
         j,
         width,
         height,
-        hdr
+        hdr,
+        denoiseAlpha
       );
+      ENGINE.endScope();
       // }, true);
       const output = outputImageData || {
         data: resGPUBuffer,
@@ -713,9 +765,9 @@ class UNet {
         height
       };
       progress?.(
-        output as T,
+        output as any,
         // Is undefined if using webgpu buffer
-        outputTileData as T | undefined,
+        outputTileData as any,
         new Tile(i * tileWidth, j * tileHeight, tileWidth, tileHeight),
         i + j * tileCountW,
         tileCountW * tileCountH
@@ -730,7 +782,8 @@ class UNet {
           }
         });
       } else {
-        done(output as T);
+        // console.log(memory());
+        done(output as any);
       }
     };
 
@@ -744,6 +797,7 @@ class UNet {
   dispose() {
     this._tfModel?.dispose();
     this._dataProcessGPU?.dispose();
+    this._tensors.forEach((tensor) => tensor.dispose());
   }
 }