oidn-web 0.2.2 → 0.3.0

package/src/process.ts

@@ -1,3 +1,4 @@
+import { buffer, data } from '@tensorflow/tfjs';
 import { WGPUComputePass } from './WGPUComputePass';
 
 const a = 1.41283765e3;
@@ -140,12 +141,53 @@ const rcpNormScale = ${rcpNormScale};
 `;
 
 export class GPUDataProcess {
-  private _inputPass;
+  private _inputPassAux;
+  private _inputPassColor;
   private _outputPass;
-  constructor(private _device: GPUDevice) {
-    this._inputPass = new WGPUComputePass('inputPass', this._device, {
+  private _copyPass;
+
+  private _isInputTexture = false;
+
+  constructor(private _device: GPUDevice, private _isHDR: boolean) {
+    const commonUniforms = [
+      {
+        label: 'inputScale',
+        type: 'f32',
+        data: new Float32Array([1])
+      },
+      {
+        label: 'inputSize',
+        type: 'vec2<f32>',
+        data: new Float32Array(2)
+      },
+      {
+        label: 'outputSize',
+        type: 'vec2<f32>',
+        data: new Float32Array(2)
+      },
+      {
+        label: 'inputOffset',
+        type: 'vec2<f32>',
+        data: new Float32Array(2)
+      }
+    ];
+    this._inputPassAux = new WGPUComputePass('inputPassAux', this._device, {
       inputs: ['color', 'albedo', 'normal'],
       outputs: ['color', 'albedo', 'normal'],
+      uniforms: commonUniforms,
+      csDefine: '',
+      csMain: ``
+    });
+    this._inputPassColor = new WGPUComputePass('inputPassColor', this._device, {
+      inputs: ['color'],
+      outputs: ['color'],
+      uniforms: commonUniforms,
+      csDefine: '',
+      csMain: ``
+    });
+    this._outputPass = new WGPUComputePass('outputPass', this._device, {
+      inputs: ['color', 'raw'],
+      outputs: ['color'],
       uniforms: [
         {
           label: 'inputScale',
@@ -162,13 +204,62 @@ export class GPUDataProcess {
           type: 'vec2<f32>',
           data: new Float32Array(2)
         },
+        {
+          label: 'imageSize',
+          type: 'vec2<f32>',
+          data: new Float32Array(2)
+        },
         {
           label: 'inputOffset',
           type: 'vec2<f32>',
           data: new Float32Array(2)
+        },
+        {
+          label: 'outputOffset',
+          type: 'vec2<f32>',
+          data: new Float32Array(2)
         }
       ],
-      csDefine: /* wgsl */ `
+      csDefine: '',
+      csMain: ``
+    });
+    this._copyPass = new WGPUComputePass('copyPass', this._device, {
+      inputs: ['color'],
+      outputs: ['color'],
+      autoUpdateOutputBuffer: false,
+      uniforms: [
+        {
+          label: 'size',
+          type: 'vec2<f32>',
+          data: new Float32Array(2)
+        }
+      ],
+      csMain: /*wgsl*/ `
+let outIdx = i32(globalId.x + globalId.y * u32(size.x));
+out_color[outIdx] = textureLoad(in_color, globalId.xy, 0);
+`
+    });
+    this._inputPassAux.setOutputParams({
+      color: { channels: 3 },
+      albedo: { channels: 3 },
+      normal: { channels: 3 }
+    });
+    this._inputPassColor.setOutputParams({
+      color: { channels: 3 }
+    });
+    this._outputPass.setOutputParams({
+      color: { channels: 4 }
+    });
+  }
+
+  private _updatePasses(isInputTexture: boolean, denoiseAlpha = false) {
+    if (this._isInputTexture === isInputTexture) {
+      return;
+    }
+
+    this._isInputTexture = isInputTexture;
+    const isHDR = this._isHDR;
+    const commonCSDefine = /* wgsl */ `
 ${constsCode}
 fn PUForward(y: f32) -> f32 {
   if (y <= y0) {
@@ -178,59 +269,49 @@ fn PUForward(y: f32) -> f32 {
   } else {
     return e * log(y + f) + g;
   }
-}
-`,
-      csMain: /* wgsl */ `
+}`;
+    function readInputCode(inputName: string) {
+      return isInputTexture
+        ? `textureLoad(in_${inputName}, globalId.xy + vec2u(inputOffset), 0)`
+        : `in_${inputName}[inIdx]'`;
+    }
+    const commonCSMain = /* wgsl */ `
 let x = f32(globalId.x);
 let y = f32(globalId.y);
 let inIdx = i32((y + inputOffset.y) * inputSize.x + (x + inputOffset.x));
-let col = in_color[inIdx] * inputScale;
-let alb = in_albedo[inIdx];
-let nor = in_normal[inIdx];
+let col = ${readInputCode('color')};
 
 let outIdx = i32(y * outputSize.x + x);
-out_color[outIdx] = vec3f(PUForward(col.r), PUForward(col.g), PUForward(col.b)) * normScale;
+
+if (${denoiseAlpha}) {
+  // Denoise the inversed alpha. Or the anti aliased edge will be too dark after denoised
+  out_color[outIdx] = vec3f(1.0 - col.a);
+}
+else if (${isHDR}) {
+  out_color[outIdx] = vec3f(PUForward(col.r * inputScale), PUForward(col.g * inputScale), PUForward(col.b * inputScale)) * normScale;
+}
+else {
+  out_color[outIdx] = col.rgb;
+}
+`;
+    this._inputPassAux.setCSCode({
+      csDefine: commonCSDefine,
+      csMain: /* wgsl */ `
+${commonCSMain}
+let alb = ${readInputCode('albedo')};
+let nor = ${readInputCode('normal')};
 out_normal[outIdx] = nor.rgb;
 out_albedo[outIdx] = alb.rgb;
-`
+  `
     });
 
-    this._outputPass = new WGPUComputePass('outputPass', this._device, {
-      inputs: ['color', 'raw'],
-      outputs: ['color'],
-      uniforms: [
-        {
-          // TODO inputScale from avg log lum
-          label: 'inputScale',
-          type: 'f32',
-          data: new Float32Array([1])
-        },
-        {
-          label: 'inputSize',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
-        },
-        {
-          label: 'outputSize',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
-        },
-        {
-          label: 'imageSize',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
-        },
-        {
-          label: 'inputOffset',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
-        },
-        {
-          label: 'outputOffset',
-          type: 'vec2<f32>',
-          data: new Float32Array(2)
-        }
-      ],
+    this._inputPassColor.setCSCode({
+      csDefine: commonCSDefine,
+      csMain: /* wgsl */ `
+${commonCSMain}
+`
+    });
+    this._outputPass.setCSCode({
       csDefine: /* wgsl */ `
 ${constsCode}
 fn PUInverse(y: f32) -> f32 {
@@ -251,40 +332,46 @@ if (x >= outputSize.x || y >= outputSize.y) {
 }
 let inIdx = i32((y + inputOffset.y) * inputSize.x + x + inputOffset.x);
 let outIdx = i32((y + outputOffset.y) * imageSize.x + x + outputOffset.x);
-let col = in_color[inIdx] * rcpNormScale;
-let raw = in_raw[outIdx];
-out_color[outIdx] = vec4f(
-  vec3f(PUInverse(col.r), PUInverse(col.g), PUInverse(col.b)) / inputScale,
-  // Pick the alpha
-  raw.a
-);
+let col = in_color[inIdx];
+let raw = ${
+        isInputTexture
+          ? 'textureLoad(in_raw, globalId.xy + vec2u(outputOffset), 0)'
+          : 'in_raw[outIdx]'
+      };
+
+if (${denoiseAlpha}) {
+  out_color[outIdx] = vec4f(raw.rgb, 1.0 - col.r);
+}
+else if (${isHDR}) {
+  out_color[outIdx] = vec4f(
+    vec3f(PUInverse(col.r * rcpNormScale), PUInverse(col.g * rcpNormScale), PUInverse(col.b * rcpNormScale)) / inputScale,
+    // Pick the alpha
+    raw.a
+  );
+}
+else {
+  out_color[outIdx] = vec4f(col.rgb, raw.a);
+}
 `
     });
-
-    // TODO input scale
-    this._inputPass.setOutputParams({
-      color: { channels: 3 },
-      albedo: { channels: 3 },
-      normal: { channels: 3 }
-    });
-    // TODO input scale
-    this._outputPass.setOutputParams({
-      color: { channels: 4 }
-    });
   }
 
   setImageSize(w: number, h: number) {
-    this._inputPass.setUniform('inputSize', new Float32Array([w, h]));
+    this._inputPassAux.setUniform('inputSize', new Float32Array([w, h]));
+    this._inputPassColor.setUniform('inputSize', new Float32Array([w, h]));
     this._outputPass.setUniform('imageSize', new Float32Array([w, h]));
     this._outputPass.setSize(w, h);
+    this._copyPass.setSize(w, h);
+    this._copyPass.setUniform('size', new Float32Array([w, h]));
   }
 
   setInputTile(tile: Tile) {
-    const inputPass = this._inputPass;
     const size = new Float32Array([tile.width, tile.height]);
-    inputPass.setUniform('inputOffset', new Float32Array([tile.x, tile.y]));
-    inputPass.setUniform('outputSize', size);
-    inputPass.setSize(size[0], size[1]);
+    [this._inputPassAux, this._inputPassColor].forEach((inputPass) => {
+      inputPass.setUniform('inputOffset', new Float32Array([tile.x, tile.y]));
+      inputPass.setUniform('outputSize', size);
+      inputPass.setSize(size[0], size[1]);
+    });
 
     this._outputPass.setUniform('inputSize', size);
   }
@@ -304,28 +391,56 @@ out_color[outIdx] = vec4f(
   }
 
   forward(
-    colorBuffer: GPUBuffer,
+    colorBuffer: GPUBuffer | GPUTexture,
     // TODO optional albedo and normal.
-    albedoBuffer: GPUBuffer,
-    normalBuffer: GPUBuffer
+    albedoBuffer: GPUBuffer | GPUTexture | undefined,
+    normalBuffer: GPUBuffer | GPUTexture | undefined,
+
+    denoiseAlpha: boolean | undefined
   ) {
-    const inputGPUPass = this._inputPass;
+    const isInputTexture = colorBuffer instanceof GPUTexture;
+    this._updatePasses(isInputTexture, denoiseAlpha);
+
+    const inputPassAux = this._inputPassAux;
+    const inputPassColor = this._inputPassColor;
     const commandEncoder = this._device.createCommandEncoder();
-    inputGPUPass.createPass(commandEncoder, {
-      color: { buffer: colorBuffer, channels: 4 },
-      albedo: { buffer: albedoBuffer, channels: 4 },
-      normal: { buffer: normalBuffer, channels: 4 }
-    });
+
+    function createInput(bufferOrTex: GPUBuffer | GPUTexture) {
+      return bufferOrTex instanceof GPUTexture
+        ? {
+            texture: bufferOrTex,
+            channels: 4
+          }
+        : {
+            buffer: bufferOrTex,
+            channels: 4
+          };
+    }
+    if (albedoBuffer && normalBuffer) {
+      inputPassAux.createPass(commandEncoder, {
+        color: createInput(colorBuffer),
+        albedo: createInput(albedoBuffer),
+        normal: createInput(normalBuffer)
+      });
+    } else {
+      inputPassColor.createPass(commandEncoder, {
+        color: createInput(colorBuffer)
+      });
+    }
     this._device.queue.submit([commandEncoder.finish()]);
 
-    return {
-      color: inputGPUPass.getOutputBuffer('color'),
-      albedo: inputGPUPass.getOutputBuffer('albedo'),
-      normal: inputGPUPass.getOutputBuffer('normal')
-    };
+    return albedoBuffer && normalBuffer
+      ? {
+          color: inputPassAux.getOutput('color'),
+          albedo: inputPassAux.getOutput('albedo'),
+          normal: inputPassAux.getOutput('normal')
+        }
+      : {
+          color: inputPassColor.getOutput('color')
+        };
   }
 
-  inverse(buffer: GPUBuffer, raw: GPUBuffer) {
+  inverse(buffer: GPUBuffer, raw: GPUBuffer | GPUTexture) {
     const device = this._device;
 
     const commandEncoder = device.createCommandEncoder();
@@ -333,30 +448,44 @@ out_color[outIdx] = vec4f(
 
     outputGPUPass.createPass(commandEncoder, {
       color: { buffer: buffer, channels: 4 },
-      raw: { buffer: raw, channels: 4 }
+      raw:
+        raw instanceof GPUBuffer
+          ? { buffer: raw, channels: 4 }
+          : { texture: raw, channels: 4 }
     });
     this._device.queue.submit([commandEncoder.finish()]);
 
-    return outputGPUPass.getOutputBuffer('color');
+    return outputGPUPass.getOutput('color');
   }
 
-  copyInputDataToOutput(inputColorBuffer: GPUBuffer) {
+  copyInputDataToOutput(inputColorBuffer: GPUBuffer | GPUTexture) {
     const encoder = this._device.createCommandEncoder();
     const outputGPUPass = this._outputPass;
-    const colorBuffer = outputGPUPass.getOutputBuffer('color');
-    encoder.copyBufferToBuffer(
-      inputColorBuffer,
-      0,
-      colorBuffer,
-      0,
-      colorBuffer.size
-    );
+    const colorBuffer = outputGPUPass.getOutput('color');
+    const copyPass = this._copyPass;
+
+    if (inputColorBuffer instanceof GPUTexture) {
+      copyPass.setOutputBuffers({
+        color: colorBuffer
+      });
+      copyPass.createPass(encoder, {
+        color: { texture: inputColorBuffer, channels: 4 }
+      });
+    } else {
+      encoder.copyBufferToBuffer(
+        inputColorBuffer,
+        0,
+        colorBuffer,
+        0,
+        colorBuffer.size
+      );
+    }
 
     this._device.queue.submit([encoder.finish()]);
   }
 
   dispose() {
     this._outputPass.dispose();
-    this._inputPass.dispose();
+    this._inputPassAux.dispose();
   }
 }