oidn-webgpu 0.0.0-stage → 0.6.0

package/src/hdr.ts

@@ -0,0 +1,398 @@
+// https://github.com/bbbbx/hdr.js/blob/main/src/hdr.ts
+function frexp(v: number): { f: number; e: number } {
+  v = Number(v);
+  const result = { f: v, e: 0 };
+
+  if (v !== 0 && Number.isFinite(v)) {
+    const absV = Math.abs(v);
+    const log2 =
+      Math.log2 ||
+      function log2(n) {
+        return Math.log(n) * Math.LOG2E;
+      };
+
+    // Math.pow(2, -exp) === Infinity when exp <= -1024
+    let e = Math.max(-1023, Math.floor(log2(absV)) + 1);
+    let f = absV * Math.pow(2.0, -e);
+
+    while (f >= 1.0) {
+      f *= 0.5;
+      e++;
+    }
+    while (f < 0.5) {
+      f *= 2;
+      e--;
+    }
+
+    if (v < 0) {
+      f = -f;
+    }
+
+    result.f = f;
+    result.e = e;
+  }
+
+  return result;
+}
+
+function ldexp(f: number, e: number) {
+  return f * Math.pow(2.0, e);
+}
+
+/**
+ * Convert 4 byte uint8 buffer to 3 channels float data
+ * @param rgbe input uint8 buffer
+ * @param float output float data
+ */
+function rgbe2float(rgbe: Uint8Array, float: Float32Array) {
+  if (rgbe[3] !== 0) {
+    const f1 = ldexp(1.0, rgbe[3] - (128 + 8));
+    float[0] = rgbe[0] * f1;
+    float[1] = rgbe[1] * f1;
+    float[2] = rgbe[2] * f1;
+  } else {
+    float[0] = float[1] = float[2] = 0.0;
+  }
+}
+
+/**
+ * Convert 3 channels float data to 4 byte uint8 buffer
+ * @param float input float data
+ * @param rgbe output uint8 buffer
+ */
+function float2rgbe(float: Float32Array, rgbe: Uint8Array) {
+  const red = float[0];
+  const green = float[1];
+  const blue = float[2];
+  const v = Math.max(red, Math.max(green, blue));
+  if (v < 1e-32) {
+    rgbe[0] = rgbe[1] = rgbe[2] = rgbe[3] = 0;
+  } else {
+    const { f, e } = frexp(v);
+    const s = (f / v) * 256.0;
+    rgbe[0] = red * s;
+    rgbe[1] = green * s;
+    rgbe[2] = blue * s;
+    rgbe[3] = e + 128;
+  }
+}
+
+/**
+ * Write float data to RGBE(.hdr) file buffer
+ * @param w image width
+ * @param h image height
+ * @param data float data, RGB 3 channels.
+ * @returns file buffer
+ */
+function write_hdr(w: number, h: number, data: Float32Array): Uint8Array {
+  const s: number[] = [];
+  const comp = 3;
+  write_hdr_core(s, w, h, comp, data);
+  return new Uint8Array(s);
+}
+
+function write_hdr_core(
+  s: number[],
+  x: number,
+  y: number,
+  comp: number,
+  data: Float32Array
+): void {
+  if (y <= 0 || x <= 0 || !data) {
+    return;
+  }
+
+  const header =
+    '#?RADIANCE\n' +
+    'FORMAT=32-bit_rle_rgbe\n' +
+    'EXPOSURE=1.0\n' +
+    '\n' +
+    '-Y ' +
+    y +
+    ' +X ' +
+    x +
+    '\n';
+  header.split('').forEach((c) => {
+    const charCode = c.charCodeAt(0);
+    s.push(charCode);
+  });
+
+  const scratch = new Uint8Array(x * 4);
+  const stbi__flip_vertically_on_write = false;
+  for (let i = 0; i < y; i++) {
+    write_hdr_scanline(
+      s,
+      x,
+      comp,
+      scratch,
+      data.subarray(comp * x * (stbi__flip_vertically_on_write ? y - 1 - i : i))
+    );
+  }
+}
+
+function write_hdr_scanline(
+  s: number[],
+  width: number,
+  ncomp: number,
+  scratch: Uint8Array,
+  scanline: Float32Array
+): void {
+  const scanlineheader = new Uint8Array([2, 2, 0, 0]);
+  const rgbe = new Uint8Array(4);
+  const linear = new Float32Array(3);
+  let x;
+
+  scanlineheader[2] = (width & 0xff00) >> 8;
+  scanlineheader[3] = width & 0x00ff;
+
+  /* skip RLE for images too small or large */
+  if (width < 8 || width >= 32768) {
+    for (x = 0; x < width; x++) {
+      switch (ncomp) {
+        case 4: /* fallthrough */
+        case 3:
+          linear[2] = scanline[x * ncomp + 2];
+          linear[1] = scanline[x * ncomp + 1];
+          linear[0] = scanline[x * ncomp + 0];
+          break;
+        default:
+          linear[0] = linear[1] = linear[2] = scanline[x * ncomp + 0];
+          break;
+      }
+      float2rgbe(linear, rgbe);
+      s.push(rgbe[0], rgbe[1], rgbe[2], rgbe[3]);
+    }
+  } else {
+    let c, r;
+    /* encode into scratch buffer */
+    for (x = 0; x < width; x++) {
+      switch (ncomp) {
+        case 4: /* fallthrough */
+        case 3:
+          linear[2] = scanline[x * ncomp + 2];
+          linear[1] = scanline[x * ncomp + 1];
+          linear[0] = scanline[x * ncomp + 0];
+          break;
+        default:
+          linear[0] = linear[1] = linear[2] = scanline[x * ncomp + 0];
+          break;
+      }
+      float2rgbe(linear, rgbe);
+      scratch[x + width * 0] = rgbe[0];
+      scratch[x + width * 1] = rgbe[1];
+      scratch[x + width * 2] = rgbe[2];
+      scratch[x + width * 3] = rgbe[3];
+    }
+
+    s.push(
+      scanlineheader[0],
+      scanlineheader[1],
+      scanlineheader[2],
+      scanlineheader[3]
+    );
+
+    /* RLE each component separately */
+    for (c = 0; c < 4; c++) {
+      const comp = scratch.subarray(width * c);
+
+      x = 0;
+      while (x < width) {
+        // find first run
+        r = x;
+        while (r + 2 < width) {
+          if (comp[r] === comp[r + 1] && comp[r] === comp[r + 2]) break;
+          ++r;
+        }
+        if (r + 2 >= width) r = width;
+        // dump up to first run
+        while (x < r) {
+          let len = r - x;
+          if (len > 128) len = 128;
+          write_dump_data(s, len, comp.subarray(x));
+          x += len;
+        }
+        // if there's a run, output it
+        if (r + 2 < width) {
+          // same test as what we break out of in search loop, so only true if we break'd
+          // find next byte after run
+          while (r < width && comp[r] == comp[x]) ++r;
+          // output run up to r
+          while (x < r) {
+            let len = r - x;
+            if (len > 127) len = 127;
+            write_run_data(s, len, comp[x]);
+            x += len;
+          }
+        }
+      }
+    }
+  }
+}
+
+function write_dump_data(s: number[], length: number, data: Uint8Array): void {
+  const lengthbyte = length & 0xff;
+  if (!(length <= 128)) throw new Error('length is greater than 128');
+  s.push(lengthbyte);
+  for (let i = 0; i < length; i++) {
+    s.push(data[i]);
+  }
+}
+
+function write_run_data(s: number[], length: number, databyte: number): void {
+  const lengthbyte = (length + 128) & 0xff;
+  if (!(length + 128 <= 255)) throw new Error('length is greater than 128');
+  s.push(lengthbyte, databyte);
+}
+
+function testMagic(uint8: Uint8Array, signature: string): boolean {
+  const l = signature.length;
+  for (let i = 0; i < l; i++) {
+    if (uint8[i] !== signature.charCodeAt(i)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+function isHdr(uint8: Uint8Array): boolean {
+  let r = testMagic(uint8, '#?RADIANCE\n');
+  if (!r) {
+    r = testMagic(uint8, '#?RGBE\n');
+  }
+  return r;
+}
+
+const MAX_HEADER_LENGTH = 1024 * 10;
+const MAX_DIMENSIONS = 1 << 24;
+
+/**
+ * Read a RGBE(.hdr) file from buffer
+ * @param uint8 RGBE(.hdr) file buffer
+ * @returns Failure reason or resolved data.
+ */
+function read_hdr(uint8: Uint8Array):
+  | string
+  | {
+      rgbFloat: Float32Array;
+      width: number;
+      height: number;
+    } {
+  let header = '';
+  let pos = 0;
+
+  if (!isHdr(uint8)) {
+    return 'Corrupt HDR image.';
+  }
+
+  // read header
+  while (!header.match(/\n\n[^\n]+\n/g) && pos < MAX_HEADER_LENGTH) {
+    header += String.fromCharCode(uint8[pos++]);
+  }
+
+  // check format
+  const format = header.match(/FORMAT=(.*)$/m)?.[1];
+  if (format !== '32-bit_rle_rgbe') {
+    return 'Unsupported HDR format: ' + format;
+  }
+
+  // parse resolution
+  const rez: string[] = header.split(/\n/).reverse()[1].split(' ');
+  if (rez[0] !== '-Y' || rez[2] !== '+X') {
+    return 'Unsupported HDR format';
+  }
+  const width = Number.parseFloat(rez[3]);
+  const height = Number.parseFloat(rez[1]);
+  if (width > MAX_DIMENSIONS || height > MAX_DIMENSIONS) {
+    return 'Very large image (corrupt?)';
+  }
+
+  let i, j;
+  let c1: number = uint8[pos];
+  let c2: number = uint8[pos + 1];
+  let len: number = uint8[pos + 2];
+
+  // not run-length encoded, so we have to actually use THIS data as a decoded
+  // pixel (note this can't be a valid pixel--one of RGB must be >= 128)
+  const notRLE: boolean = c1 !== 2 || c2 !== 2 || !!(len & 0x80); // not run-length encoded
+
+  const hdrData = new Float32Array(width * height * 3);
+  if (width < 8 || width >= 32768 || notRLE) {
+    // 32768: 2^15
+    // Read flat data
+    for (j = 0; j < height; ++j) {
+      for (i = 0; i < width; ++i) {
+        const rgbe = uint8.subarray(pos, pos + 4);
+        pos += 4;
+        const start = (j * width + i) * 3;
+        rgbe2float(rgbe, hdrData.subarray(start, start + 3));
+      }
+    }
+  } else {
+    // Read RLE-encoded data
+    let scanline: Uint8Array | undefined;
+    let c1: number;
+    let c2: number;
+    let len: number;
+    for (let j = 0; j < height; j++) {
+      c1 = uint8[pos++];
+      c2 = uint8[pos++];
+      len = uint8[pos++];
+      if (c1 !== 2 || c2 !== 2 || len & 0x80) {
+        return 'Invalid scanline';
+      }
+
+      len = len << 8;
+      len |= uint8[pos++];
+      if (len !== width) {
+        return 'invalid decoded scanline length';
+      }
+      if (!scanline) {
+        scanline = new Uint8Array(width * 4);
+      }
+
+      let count: number;
+      let value: number;
+      for (let k = 0; k < 4; k++) {
+        let nLeft: number;
+        i = 0;
+        while ((nLeft = width - i) > 0) {
+          count = uint8[pos++];
+          if (count > 128) {
+            // is RUN
+            value = uint8[pos++];
+            count -= 128;
+            if (count > nLeft) {
+              return 'bad RLE data in HDR';
+            }
+            for (let z = 0; z < count; z++) {
+              scanline[i++ * 4 + k] = value;
+            }
+          } else {
+            // is DUMP
+            if (count > nLeft) {
+              return 'bad RLE data in HDR';
+            }
+            for (let z = 0; z < count; z++) {
+              scanline[i++ * 4 + k] = uint8[pos++];
+            }
+          }
+        }
+      }
+
+      for (let i = 0; i < width; i++) {
+        rgbe2float(
+          scanline.subarray(i * 4),
+          hdrData.subarray((j * width + i) * 3)
+        );
+      }
+    }
+  }
+
+  return {
+    rgbFloat: hdrData,
+    width: width,
+    height: height
+  };
+}
+
+export { read_hdr as readHDR, write_hdr as writeHDR, float2rgbe, rgbe2float };

package/src/hdrTransfer.ts

@@ -0,0 +1,169 @@
+import { FLT_MAX, sanitize } from './sanitize.js';
+
+/** HDR transfer functions used by the OIDN input and output processing passes. */
+export type HDRTransfer = 'pu' | 'log';
+
+const a = 1.41283765e3;
+const b = 1.64593172;
+const c = 4.31384981e-1;
+const d = -2.94139609e-3;
+const e = 1.92653254e-1;
+const f = 6.26026094e-3;
+const g = 9.98620152e-1;
+const y0 = 1.5794576e-6;
+const y1 = 3.22087631e-2;
+const x0 = 2.23151711e-3;
+const x1 = 3.70974749e-1;
+const yMax = 65504;
+const puXMax = puForward(yMax);
+const puNormScale = 1 / puXMax;
+const puRcpNormScale = puXMax;
+const logXMax = Math.log(yMax + 1);
+const logNormScale = 1 / logXMax;
+
+function puForward(y: number) {
+  if (y <= y0) return a * y;
+  if (y <= y1) return b * Math.pow(y, c) + d;
+  return e * Math.log(y + f) + g;
+}
+
+function puInverse(x: number) {
+  if (x <= x0) return x / a;
+  if (x <= x1) return Math.pow((x - d) / b, 1 / c);
+  return Math.exp((x - g) / e) - f;
+}
+
+function forward(y: number, transfer: HDRTransfer) {
+  return transfer === 'log'
+    ? Math.log(y + 1) * logNormScale
+    : puForward(y) * puNormScale;
+}
+
+function inverse(x: number, transfer: HDRTransfer) {
+  return transfer === 'log'
+    ? Math.exp(x * logXMax) - 1
+    : puInverse(x * puRcpNormScale);
+}
+
+export function hdrTransferFuncCPU({
+  data,
+  channels,
+  inputScale,
+  transfer = 'pu'
+}: {
+  data: Float32Array;
+  channels: number;
+  inputScale: number;
+  transfer?: HDRTransfer;
+}) {
+  const newData = new Float32Array(data);
+  for (let i = 0; i < newData.length; i += channels) {
+    for (let channel = 0; channel < 3; channel++) {
+      newData[i + channel] = forward(
+        sanitize(newData[i + channel] * inputScale, 0, FLT_MAX),
+        transfer
+      );
+    }
+  }
+  return newData;
+}
+
+export function hdrTransferFuncInverseCPU({
+  data,
+  channels,
+  inputScale,
+  transfer = 'pu'
+}: {
+  data: Float32Array;
+  channels: number;
+  inputScale: number;
+  transfer?: HDRTransfer;
+}) {
+  const newData = new Float32Array(data);
+  const outputScale = 1 / inputScale;
+  for (let i = 0; i < newData.length; i += channels) {
+    for (let channel = 0; channel < 3; channel++) {
+      // The network output may contain negative values or NaNs.
+      newData[i + channel] = inverse(sanitize(newData[i + channel], 0, FLT_MAX), transfer) * outputScale;
+    }
+  }
+  return newData;
+}
+
+// TODO put in gpu
+export function avgLogLum({
+  data,
+  channels
+}: {
+  data: Float32Array;
+  channels: number;
+}) {
+  let totalLuminance = 0;
+
+  for (let i = 0; i < data.length; i += channels) {
+    const r = sanitize(data[i], 0, FLT_MAX);
+    const g = sanitize(data[i + 1], 0, FLT_MAX);
+    const b = sanitize(data[i + 2], 0, FLT_MAX);
+    const lum = 0.212671 * r + 0.71516 * g + 0.072169 * b;
+
+    // TODO performance
+    // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/devices/cpu/cpu_autoexposure.cpp#L43
+    totalLuminance += Math.log2(lum + 0.0001);
+  }
+  const numPixels = data.length / channels;
+  const averageLuminance = totalLuminance / numPixels;
+  const key = 0.18;
+
+  // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/training/color.py#L173
+  return key / Math.pow(2, averageLuminance);
+}
+
+const transferConstantsWGSL = `
+const a = ${a};
+const b = ${b};
+const c = ${c};
+const d = ${d};
+const e = ${e};
+const f = ${f};
+const g = ${g};
+const y0 =${y0};
+const y1 =${y1};
+const x0 =${x0};
+const x1 =${x1};
+
+const normScale = ${puNormScale};
+const rcpNormScale = ${puRcpNormScale};
+`;
+
+/** WGSL declarations of `fn HDRForward(y: f32) -> f32` for the input pass. */
+export function hdrForwardWGSL(transfer: HDRTransfer) {
+  const forwardCode = transfer === 'log'
+    ? `fn HDRForward(y: f32) -> f32 { return log(y + 1.0) * logNormScale; }`
+    : `fn HDRForward(y: f32) -> f32 {
+  if (y <= y0) { return a * y * normScale; }
+  else if (y <= y1) { return (b * pow(y, c) + d) * normScale; }
+  else { return (e * log(y + f) + g) * normScale; }
+}`;
+  return /* wgsl */ `
+${transferConstantsWGSL}
+const logXMax = ${logXMax};
+const logNormScale = ${logNormScale};
+${forwardCode}`;
+}
+
+/** WGSL declarations of `fn HDRInverse(x: f32) -> f32` for the output pass. */
+export function hdrInverseWGSL(transfer: HDRTransfer) {
+  const inverseCode = transfer === 'log'
+    ? `fn HDRInverse(x: f32) -> f32 { return exp(x * logXMax) - 1.0; }`
+    : `fn HDRInverse(x: f32) -> f32 {
+  let y = x * rcpNormScale;
+  if (y <= x0) { return y / a; }
+  else if (y <= x1) { return pow((y - d) / b, 1 / c); }
+  else { return exp((y - g) / e) - f; }
+}`;
+  return /* wgsl */ `
+${transferConstantsWGSL}
+const logXMax = ${logXMax};
+${inverseCode}
+`;
+}

package/src/main.ts

@@ -0,0 +1,126 @@
+import { parseTZA } from './tza';
+import UNet from './UNet';
+import type { UNetEngineSetting, UNetExecutionStats } from './UNet';
+import { initWebGPUBackend } from './backend';
+import type { DynamicTileSetting } from './tileScheduler';
+import type { UNetModelSpec } from './modelSpec';
+import type {
+  NativeSubgroupMatrixOptions,
+  NativeUNetGemmOptions,
+  NativeUNetKernelSetting,
+  NativeUNetPrecisionSetting
+} from './nativeUNet';
+import type { HDRTransfer } from './process';
+
+export { parseTZA, UNet };
+export { planTileGrid } from './tileScheduler';
+export type {
+  DynamicTileOptions,
+  DynamicTileSetting,
+  PlannedTile,
+  TilePlan,
+  TileRect
+} from './tileScheduler';
+export {
+  detectUNetModelSpec,
+  OIDN_UNET_LARGE_SPEC,
+  OIDN_UNET_SMALL_SPEC,
+  validateUNetModel
+} from './modelSpec';
+export type {
+  ModelNodeSpec,
+  UNetModelGraph,
+  UNetModelSpec,
+  ValidatedUNetModel
+} from './modelSpec';
+export { optimizeModelGraph, planModelExecution } from './graphOptimizer';
+export type {
+  ExecutableModelNode,
+  GraphOptimizationOptions,
+  ModelExecutionPlan,
+  ModelValueShape,
+  OptimizedModelGraph
+} from './graphOptimizer';
+export {
+  NativeUNetExecutor,
+  resolveNativeUNetPrecision,
+  selectSubgroupMatrixConfig
+} from './nativeUNet';
+export type {
+  NativeSubgroupMatrixConfig,
+  NativeSubgroupMatrixOptions,
+  NativeUNetGemmWorkgroup,
+  NativeUNetGemmOptions,
+  NativeUNetExecutionProfile,
+  NativeUNetKernel,
+  NativeUNetKernelSetting,
+  NativeUNetLayerTiming,
+  NativeUNetOptions,
+  NativeUNetPrecision,
+  NativeUNetPrecisionSetting
+} from './nativeUNet';
+
+export interface UNetOptions {
+  aux?: boolean;
+  hdr?: boolean;
+  /** HDR transfer function expected by the trained model. Defaults to PU. */
+  hdrTransfer?: HDRTransfer;
+  /** Hard upper bound for an output tile edge. Defaults to 512. */
+  maxTileSize?: number;
+  /** Adaptive GPU-time-based tile sizing. Enabled by default. */
+  dynamicTile?: DynamicTileSetting;
+  /** `auto` uses native WGSL; `webnn` opts into the WebNN backend. */
+  engine?: UNetEngineSetting;
+  /** `auto` selects FP16 when shader-f16 was enabled on the GPUDevice. */
+  precision?: NativeUNetPrecisionSetting;
+  /** `auto` uses implicit GEMM for FP16/FP32 convolutions, except the direct output layer. */
+  kernel?: NativeUNetKernelSetting;
+  /**
+   * Experimental implicit-GEMM tuning for native WGSL execution. These knobs
+   * exist for benchmarking and may change or be removed in any release.
+   * @experimental
+   */
+  gemm?: NativeUNetGemmOptions;
+  /**
+   * Experimental `kernel: 'subgroup-matrix'` tuning. The kernel needs a
+   * GPUDevice created with `subgroups` and `chromium-experimental-subgroup-matrix`.
+   * @experimental
+   */
+  subgroupMatrix?: NativeSubgroupMatrixOptions;
+  /** Versioned topology descriptor for future/custom OIDN TZA models. */
+  modelSpec?: UNetModelSpec;
+}
+
+export type { UNetEngineSetting, UNetExecutionStats } from './UNet';
+export type { HDRTransfer } from './process';
+export { WebNNUNetExecutor } from './webnnUNet';
+export type { WebNNRuntimeSupport, WebNNUNetOptions } from './webnnUNet';
+export type {
+  OIDNResourceKind,
+  OIDNResourceSnapshot,
+  OIDNResourceStats
+} from './resourceTracker';
+
+export async function initUNetFromBuffer(
+  tzaBuffer: ArrayBuffer,
+  backendParams?: { device: GPUDevice },
+  opts?: UNetOptions
+) {
+  const device = backendParams?.device ?? await initWebGPUBackend();
+  const tensors = parseTZA(tzaBuffer);
+  const unet = new UNet(tensors, device, opts);
+  await unet.prepare();
+  return unet;
+}
+
+export async function initUNetFromURL(
+  modelPath: string,
+  backendParams?: { device: GPUDevice },
+  opts?: UNetOptions
+) {
+  return fetch(modelPath)
+    .then((res) => res.arrayBuffer())
+    .then((ab) => {
+      return initUNetFromBuffer(ab, backendParams, opts);
+    });
+}