oidn-web 0.4.0 → 0.5.0

package/src/process.ts

@@ -1,4 +1,11 @@
 import { WGPUComputePass } from './WGPUComputePass';
+import {
+  hdrTransferFuncCPU,
+  hdrTransferFuncInverseCPU,
+  type HDRTransfer
+} from './hdrTransfer';
+export { hdrTransferFuncCPU, hdrTransferFuncInverseCPU } from './hdrTransfer';
+export type { HDRTransfer } from './hdrTransfer';
 
 const a = 1.41283765e3;
 const b = 1.64593172;
@@ -37,6 +44,8 @@ const yMax = 65504;
 const xMax = PUForward(yMax);
 const normScale = 1 / xMax;
 const rcpNormScale = xMax;
+const logXMax = Math.log(yMax + 1);
+const logNormScale = 1 / logXMax;
 
 export class Tile {
   constructor(
@@ -75,53 +84,6 @@ export function avgLogLum({
   return key / Math.pow(2, averageLuminance);
 }
 
-export function hdrTransferFuncCPU({
-  data,
-  channels,
-  inputScale
-}: {
-  data: Float32Array;
-  channels: number;
-  inputScale: number;
-}) {
-  const newData = new Float32Array(data.length);
-  newData.set(data);
-  // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/core/color.h#L71
-  for (let i = 0; i < newData.length; i += channels) {
-    // First three are color
-    for (let c = 0; c < 3; c++) {
-      let y = newData[i + c] * inputScale;
-      // https://github.com/RenderKit/oidn/blob/713ec7838ba650f99e0a896549c0dca5eeb3652d/devices/cpu/color.ispc#L135
-      newData[i + c] = PUForward(y) * normScale;
-    }
-  }
-
-  return newData;
-}
-
-export function hdrTransferFuncInverseCPU({
-  data,
-  channels,
-  inputScale
-}: {
-  data: Float32Array;
-  channels: number;
-  inputScale: number;
-}) {
-  const newData = new Float32Array(data.length);
-  newData.set(data);
-
-  const outputScale = 1 / inputScale;
-  for (let i = 0; i < newData.length; i += channels) {
-    for (let c = 0; c < 3; c++) {
-      let x = newData[i + c] * rcpNormScale;
-      newData[i + c] = PUInverse(x) * outputScale;
-    }
-  }
-
-  return newData;
-}
-
 const constsCode = `
 const a = ${a};
 const b = ${b};
@@ -147,7 +109,11 @@ export class GPUDataProcess {
 
   private _isInputTexture?: boolean;
 
-  constructor(private _device: GPUDevice, private _isHDR: boolean) {
+  constructor(
+    private _device: GPUDevice,
+    private _isHDR: boolean,
+    private _hdrTransfer: HDRTransfer = 'pu'
+  ) {
     const commonUniforms = [
       {
         label: 'inputScale',
@@ -261,26 +227,40 @@ out_color[outIdx] = textureLoad(in_color, globalId.xy, 0);
 
     this._isInputTexture = isInputTexture;
     const isHDR = this._isHDR;
+    const hdrForwardCode = this._hdrTransfer === '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; }
+}`;
+    const hdrInverseCode = this._hdrTransfer === '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; }
+}`;
     const commonCSDefine = /* wgsl */ `
 ${constsCode}
-fn PUForward(y: f32) -> f32 {
-  if (y <= y0) {
-    return a * y;
-  } else if (y <= y1) {
-    return b * pow(y, c) + d;
-  } else {
-    return e * log(y + f) + g;
-  }
-}`;
+const logXMax = ${logXMax};
+const logNormScale = ${logNormScale};
+${hdrForwardCode}`;
     function readInputCode(inputName: string) {
       return isInputTexture
-        ? `textureLoad(in_${inputName}, globalId.xy + vec2u(inputOffset), 0)`
+        ? `textureLoad(in_${inputName}, vec2u(inputPosition), 0)`
         : `in_${inputName}[inIdx]`;
     }
     const commonCSMain = /* wgsl */ `
 let x = i32(globalId.x);
 let y = i32(globalId.y);
-let inIdx = (y + inputOffset.y) * inputSize.x + (x + inputOffset.x);
+let inputPosition = clamp(
+  vec2i(x, y) + inputOffset,
+  vec2i(0),
+  inputSize - vec2i(1)
+);
+let inIdx = inputPosition.y * inputSize.x + inputPosition.x;
 let col = ${readInputCode('color')};
 
 let outIdx = y * outputSize.x + x;
@@ -290,7 +270,7 @@ if (${denoiseAlpha}) {
   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;
+  out_color[outIdx] = vec3f(HDRForward(col.r * inputScale), HDRForward(col.g * inputScale), HDRForward(col.b * inputScale));
 }
 else {
   out_color[outIdx] = col.rgb;
@@ -316,15 +296,8 @@ ${commonCSMain}
     this._outputPass.setCSCode({
       csDefine: /* wgsl */ `
 ${constsCode}
-fn PUInverse(y: f32) -> f32 {
-  if (y <= x0) {
-    return y / a;
-  } else if (y <= x1) {
-    return pow((y - d) / b, 1 / c);
-  } else {
-    return exp((y - g) / e) - f;
-  }
-}
+const logXMax = ${logXMax};
+${hdrInverseCode}
 `,
       csMain: /* wgsl */ `
 let x = i32(globalId.x);
@@ -346,7 +319,7 @@ if (${denoiseAlpha}) {
 }
 else if (${isHDR}) {
   out_color[outIdx] = vec4f(
-    vec3f(PUInverse(col.r * rcpNormScale), PUInverse(col.g * rcpNormScale), PUInverse(col.b * rcpNormScale)) / inputScale,
+    vec3f(HDRInverse(col.r), HDRInverse(col.g), HDRInverse(col.b)) / inputScale,
     // Pick the alpha
     raw.a
   );

package/src/tileScheduler.ts

@@ -1,22 +1,47 @@
 export interface DynamicTileOptions {
   /** Smallest output tile edge. Defaults to 256. */
   minTileSize?: number;
-  /** Initial output tile edge. Defaults to 384. */
+  /** Initial maximum output tile edge. Defaults to 432. */
   initialTileSize?: number;
   /** Desired upper bound for one tile's GPU latency. Defaults to 16 ms. */
   targetTileTimeMs?: number;
-  /** Amount added or removed after each completed execution. Defaults to 128. */
+  /** Amount added or removed after each completed execution. Defaults to 16. */
   adjustmentStep?: number;
 }
 
 export type DynamicTileSetting = boolean | DynamicTileOptions;
 
 const defaultMinTileSize = 256;
-const defaultInitialTileSize = 384;
+const defaultInitialTileSize = 432;
 const defaultTargetTileTimeMs = 16;
-const defaultAdjustmentStep = 128;
+const defaultAdjustmentStep = 16;
 export const OIDN_TILE_ALIGNMENT = 16;
 
+export interface TileRect {
+  x: number;
+  y: number;
+  width: number;
+  height: number;
+}
+
+export interface PlannedTile {
+  column: number;
+  row: number;
+  input: TileRect;
+  output: TileRect;
+}
+
+export interface TilePlan {
+  columns: number;
+  rows: number;
+  overlap: number;
+  maxOutputWidth: number;
+  maxOutputHeight: number;
+  inputPixelCount: number;
+  inputShapeCount: number;
+  tiles: PlannedTile[];
+}
+
 function alignUp(value: number, alignment: number) {
   return Math.ceil(value / alignment) * alignment;
 }
@@ -25,16 +50,180 @@ function alignDown(value: number, alignment: number) {
   return Math.floor(value / alignment) * alignment;
 }
 
+function partitionRange(length: number, parts: number, index: number) {
+  const start = Math.round(index * length / parts);
+  const end = Math.round((index + 1) * length / parts);
+  return { start, end };
+}
+
+function alignRangeWithinImage(
+  start: number,
+  end: number,
+  imageSize: number,
+  minimumSize = 0
+) {
+  const currentSize = end - start;
+  const alignedSize = Math.max(
+    minimumSize,
+    alignUp(currentSize, OIDN_TILE_ALIGNMENT)
+  );
+  let alignedStart = start - Math.floor((alignedSize - currentSize) / 2);
+  // At an image edge, keep the padded portion on the right/bottom. Input
+  // readers replicate the nearest image edge into that portion, while output
+  // tiles remain strictly inside the original image.
+  alignedStart = clamp(alignedStart, 0, Math.max(0, imageSize - alignedSize));
+  return { start: alignedStart, end: alignedStart + alignedSize };
+}
+
+/**
+ * Splits an image into balanced rectangular output regions. Only sides that
+ * touch another tile receive a halo; image edges rely on the model's normal
+ * boundary behavior instead of doing redundant off-edge work.
+ */
+export function planTileGrid(
+  width: number,
+  height: number,
+  maxTileSize: number,
+  overlap: number
+): TilePlan {
+  if (!Number.isInteger(width) || width <= 0 ||
+      !Number.isInteger(height) || height <= 0) {
+    throw new Error('Tile grid dimensions must be positive integers');
+  }
+  if (!Number.isFinite(maxTileSize) || maxTileSize <= 0) {
+    throw new Error('Maximum tile size must be positive');
+  }
+  if (!Number.isFinite(overlap) || overlap < 0) {
+    throw new Error('Tile overlap must be non-negative');
+  }
+
+  const alignedMaximum = Math.max(
+    OIDN_TILE_ALIGNMENT,
+    alignDown(maxTileSize, OIDN_TILE_ALIGNMENT)
+  );
+  const alignedOverlap = alignUp(overlap, OIDN_TILE_ALIGNMENT);
+  const columns = Math.max(1, Math.ceil(width / alignedMaximum));
+  const rows = Math.max(1, Math.ceil(height / alignedMaximum));
+  const tiles: PlannedTile[] = [];
+  let maxOutputWidth = 0;
+  let maxOutputHeight = 0;
+
+  for (let row = 0; row < rows; row++) {
+    const outputY = partitionRange(height, rows, row);
+    for (let column = 0; column < columns; column++) {
+      const outputX = partitionRange(width, columns, column);
+      const desiredInputX = {
+        start: column === 0
+          ? outputX.start
+          : Math.max(0, outputX.start - alignedOverlap),
+        end: column === columns - 1
+          ? outputX.end
+          : Math.min(width, outputX.end + alignedOverlap)
+      };
+      const desiredInputY = {
+        start: row === 0
+          ? outputY.start
+          : Math.max(0, outputY.start - alignedOverlap),
+        end: row === rows - 1
+          ? outputY.end
+          : Math.min(height, outputY.end + alignedOverlap)
+      };
+      const inputX = alignRangeWithinImage(
+        desiredInputX.start,
+        desiredInputX.end,
+        width
+      );
+      const inputY = alignRangeWithinImage(
+        desiredInputY.start,
+        desiredInputY.end,
+        height
+      );
+      const output: TileRect = {
+        x: outputX.start,
+        y: outputY.start,
+        width: outputX.end - outputX.start,
+        height: outputY.end - outputY.start
+      };
+      const input: TileRect = {
+        x: inputX.start,
+        y: inputY.start,
+        width: inputX.end - inputX.start,
+        height: inputY.end - inputY.start
+      };
+      maxOutputWidth = Math.max(maxOutputWidth, output.width);
+      maxOutputHeight = Math.max(maxOutputHeight, output.height);
+      tiles.push({ column, row, input, output });
+    }
+  }
+
+  const inputShapes = () => new Set(
+    tiles.map(({ input }) => `${input.width}x${input.height}`)
+  );
+  if (inputShapes().size > 2) {
+    const maximumInputWidth = Math.max(
+      ...tiles.map(({ input }) => input.width)
+    );
+    const maximumInputHeight = Math.max(
+      ...tiles.map(({ input }) => input.height)
+    );
+    const uniformWidthPixelCount = tiles.reduce(
+      (sum, { input }) => sum + maximumInputWidth * input.height,
+      0
+    );
+    const uniformHeightPixelCount = tiles.reduce(
+      (sum, { input }) => sum + input.width * maximumInputHeight,
+      0
+    );
+    if (uniformWidthPixelCount <= uniformHeightPixelCount) {
+      for (const tile of tiles) {
+        const range = alignRangeWithinImage(
+          tile.input.x,
+          tile.input.x + tile.input.width,
+          width,
+          maximumInputWidth
+        );
+        tile.input.x = range.start;
+        tile.input.width = range.end - range.start;
+      }
+    } else {
+      for (const tile of tiles) {
+        const range = alignRangeWithinImage(
+          tile.input.y,
+          tile.input.y + tile.input.height,
+          height,
+          maximumInputHeight
+        );
+        tile.input.y = range.start;
+        tile.input.height = range.end - range.start;
+      }
+    }
+  }
+
+  const finalInputShapes = inputShapes();
+  const inputPixelCount = tiles.reduce(
+    (sum, { input }) => sum + input.width * input.height,
+    0
+  );
+
+  return {
+    columns,
+    rows,
+    overlap: alignedOverlap,
+    maxOutputWidth,
+    maxOutputHeight,
+    inputPixelCount,
+    inputShapeCount: finalInputShapes.size,
+    tiles
+  };
+}
+
 function clamp(value: number, min: number, max: number) {
   return Math.min(Math.max(value, min), max);
 }
 
-function median(values: number[]) {
+function percentile(values: number[], fraction: number) {
   const sorted = [...values].sort((a, b) => a - b);
-  const middle = Math.floor(sorted.length / 2);
-  return sorted.length % 2
-    ? sorted[middle]
-    : (sorted[middle - 1] + sorted[middle]) / 2;
+  return sorted[Math.ceil(sorted.length * fraction) - 1];
 }
 
 /** Fits one output dimension without exceeding the configured tile cap. */
@@ -57,6 +246,8 @@ export class DynamicTileController {
 
   private _tileSize: number;
   private _adjustmentStep: number;
+  private _smoothedTileTimeMs?: number;
+  private _completeExecutionsSinceChange = 0;
 
   constructor(maxTileSize: number, setting: DynamicTileSetting = true) {
     const options = typeof setting === 'object' ? setting : {};
@@ -106,11 +297,22 @@ export class DynamicTileController {
     );
     if (validTimes.length === 0) return false;
 
-    const observedTime = median(validTimes);
+    // The first tile may still pay one-time GPU residency costs even after
+    // shape preparation. It should not make a stable layout look too slow.
+    const representativeTimes = validTimes.length >= 3
+      ? validTimes.slice(1)
+      : validTimes;
+    const observedTime = percentile(representativeTimes, 0.75);
+    this._smoothedTileTimeMs = this._smoothedTileTimeMs === undefined
+      ? observedTime
+      : this._smoothedTileTimeMs * 0.65 + observedTime * 0.35;
+    this._completeExecutionsSinceChange++;
+    if (this._completeExecutionsSinceChange < 2) return false;
+
     let nextTileSize = this._tileSize;
-    if (observedTime > this.targetTileTimeMs * 1.25) {
+    if (this._smoothedTileTimeMs > this.targetTileTimeMs * 1.25) {
       nextTileSize -= this._adjustmentStep;
-    } else if (observedTime < this.targetTileTimeMs * 0.65) {
+    } else if (this._smoothedTileTimeMs < this.targetTileTimeMs * 0.65) {
       nextTileSize += this._adjustmentStep;
     }
 
@@ -121,18 +323,8 @@ export class DynamicTileController {
     );
     if (nextTileSize === this._tileSize) return false;
     this._tileSize = nextTileSize;
+    this._smoothedTileTimeMs = undefined;
+    this._completeExecutionsSinceChange = 0;
     return true;
   }
 }
-
-/** Waits for all work submitted before this call and tolerates device loss. */
-export async function waitForSubmittedGPUWork(
-  queue: Pick<GPUQueue, 'onSubmittedWorkDone'>
-) {
-  try {
-    await queue.onSubmittedWorkDone();
-  } catch {
-    // tileExecute has no error callback. Preserve its completion/cancellation
-    // behavior and let the owning GPUDevice report device loss separately.
-  }
-}