@driftengine/capture 4.0.0

package/src/models/dualDpt.ts

@@ -0,0 +1,233 @@
+/**
+ * Depth Anything 3's head, the DPT that turns four taps of tokens into depth and its confidence.
+ *
+ * **Each tap is normed, laid back out as an image, projected, and resized to its level of a
+ * pyramid** — four times the patch grid, twice, once, and half — then four fusion blocks walk the
+ * pyramid from coarse to fine, each adding the level below through a residual unit, and the result
+ * is brought up to the image by an aligned bilinear resize and read out by two convolutions: a
+ * depth logit and a confidence logit per pixel. Their activations, `exp` and `exp + 1`, are the
+ * caller's, applied to the two channels as they are read back, since the runtime has no `exp` and
+ * one pass over a map on the host costs less than a kernel for it.
+ *
+ * **A sinusoidal position is added three times**, after each projection and before the readout:
+ * a grid spanning the image's aspect normalised by its diagonal, as the upstream's `create_uv_grid`
+ * makes it, embedded at base 100 and scaled by a tenth. It depends only on shapes, so each is a
+ * constant of the graph.
+ *
+ * **The upstream's second head is not built.** It predicts rays for a pose estimate the camera
+ * decoder already makes, and the upstream deletes its output when that decoder is present, as it is
+ * in every accepted checkpoint; its weights are set aside by name.
+ */
+import { exactCos, exactExp, exactSin } from '@driftengine/core';
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+import { LOG_BASE } from './rope.ts';
+
+export interface HeadConfig {
+  readonly dimIn: number;
+  readonly features: number;
+  readonly outChannels: readonly [number, number, number, number];
+  readonly outputDim: number;
+  readonly patch: number;
+}
+
+/** The head's own norms keep PyTorch's default ε. */
+const EPSILON = 1e-5;
+
+const f = Math.fround;
+
+/**
+ * `[channels, height, width]`: the upstream's `_add_pos_embed` for a map of this size in an image of
+ * `aspect` width over height — sine and cosine of the horizontal coordinate, then of the vertical,
+ * over `channels / 4` frequencies each, times a tenth.
+ */
+export function uvEmbedding(
+  channels: number,
+  height: number,
+  width: number,
+  aspect: number,
+): Float32Array {
+  const diagonal = Math.sqrt(aspect * aspect + 1);
+  const spanX = aspect / diagonal;
+  const spanY = 1 / diagonal;
+  const coordinate = (span: number, count: number, i: number): number => {
+    const edge = (span * (count - 1)) / count;
+    return count === 1 ? -edge : f(-edge + ((2 * edge) / (count - 1)) * i);
+  };
+  const quarter = channels / 4;
+  const omega = Float32Array.from({ length: quarter }, (_, k) =>
+    f(1 / f(exactExp(f(k / quarter) * LOG_BASE))),
+  );
+  const out = new Float32Array(channels * height * width);
+  for (let y = 0; y < height; y += 1) {
+    const v = coordinate(spanY, height, y);
+    for (let x = 0; x < width; x += 1) {
+      const u = coordinate(spanX, width, x);
+      for (let k = 0; k < quarter; k += 1) {
+        const au = f(u * (omega[k] as number));
+        const av = f(v * (omega[k] as number));
+        const at = y * width + x;
+        const plane = height * width;
+        out[k * plane + at] = f(f(exactSin(au)) * 0.1);
+        out[(quarter + k) * plane + at] = f(f(exactCos(au)) * 0.1);
+        out[(2 * quarter + k) * plane + at] = f(f(exactSin(av)) * 0.1);
+        out[(3 * quarter + k) * plane + at] = f(f(exactCos(av)) * 0.1);
+      }
+    }
+  }
+  return out;
+}
+
+/** The size a stride-2 3×3 convolution with padding 1 leaves. */
+const halved = (size: number): number => Math.floor((size - 1) / 2) + 1;
+
+/**
+ * The head over one view's four taps, each `[rows · cols, dimIn]`; the result is the logits,
+ * `[outputDim, rows · patch, cols · patch]`, named `output`. Several views share its constants.
+ */
+export function dualDptHead(
+  weights: Weights,
+  graph: GraphBuilder,
+  config: HeadConfig,
+  prefix: string,
+  taps: readonly string[],
+  rows: number,
+  cols: number,
+  output: string,
+): string {
+  const node = graph.node.bind(graph);
+  const w = (name: string, shape: readonly number[]): string =>
+    weights.read(`${prefix}${name}`, shape);
+  const height = rows * config.patch;
+  const width = cols * config.patch;
+  const aspect = width / height;
+  const features = config.features;
+  const embed = (x: string, channels: number, h: number, wide: number): string =>
+    node('add', [
+      x,
+      weights.constant(
+        `uv.${channels}.${h}.${wide}`,
+        [channels, h, wide],
+        uvEmbedding(channels, h, wide, aspect),
+      ),
+    ]);
+  const conv = (
+    x: string,
+    name: string,
+    from: number,
+    to: number,
+    kernel: number,
+    attributes = {},
+    bias = true,
+    named?: string,
+  ): string =>
+    node(
+      'conv2d',
+      bias
+        ? [x, w(`${name}.weight`, [to, from, kernel, kernel]), w(`${name}.bias`, [to])]
+        : [x, w(`${name}.weight`, [to, from, kernel, kernel])],
+      attributes,
+      named,
+    );
+
+  const sizes: [number, number][] = [
+    [4 * rows, 4 * cols],
+    [2 * rows, 2 * cols],
+    [rows, cols],
+    [halved(rows), halved(cols)],
+  ];
+  const levels = taps.map((tap, k) => {
+    const channels = config.outChannels[k] as number;
+    const normed = node(
+      'layerNorm',
+      [tap, w('norm.weight', [config.dimIn]), w('norm.bias', [config.dimIn])],
+      { epsilon: EPSILON },
+    );
+    const image = node('reshape', [node('permute', [normed], { order: [1, 0] })], {
+      shape: [config.dimIn, rows, cols],
+    });
+    const projected = embed(
+      conv(image, `projects.${k}`, config.dimIn, channels, 1),
+      channels,
+      rows,
+      cols,
+    );
+    const resized =
+      k === 0
+        ? node(
+            'convTranspose2d',
+            [
+              projected,
+              w('resize_layers.0.weight', [channels, channels, 4, 4]),
+              w('resize_layers.0.bias', [channels]),
+            ],
+            { stride: 4 },
+          )
+        : k === 1
+          ? node(
+              'convTranspose2d',
+              [
+                projected,
+                w('resize_layers.1.weight', [channels, channels, 2, 2]),
+                w('resize_layers.1.bias', [channels]),
+              ],
+              { stride: 2 },
+            )
+          : k === 2
+            ? projected
+            : conv(projected, 'resize_layers.3', channels, channels, 3, { stride: 2, padding: 1 });
+    return conv(resized, `scratch.layer${k + 1}_rn`, channels, features, 3, { padding: 1 }, false);
+  });
+
+  const residual = (x: string, name: string): string => {
+    const first = conv(node('relu', [x]), `${name}.conv1`, features, features, 3, { padding: 1 });
+    return node('add', [
+      conv(node('relu', [first]), `${name}.conv2`, features, features, 3, { padding: 1 }),
+      x,
+    ]);
+  };
+  const fuse = (
+    level: number,
+    top: string,
+    lateral: string | null,
+    size: [number, number],
+  ): string => {
+    const name = `scratch.refinenet${level}`;
+    const merged =
+      lateral === null ? top : node('add', [top, residual(lateral, `${name}.resConfUnit1`)]);
+    const refined = residual(merged, `${name}.resConfUnit2`);
+    const up = node('resize', [refined], {
+      height: size[0],
+      width: size[1],
+      mode: 'bilinear',
+      alignCorners: true,
+    });
+    return conv(up, `${name}.out_conv`, features, features, 1);
+  };
+  const [l1, l2, l3, l4] = levels as [string, string, string, string];
+  let out = fuse(4, l4, null, sizes[2]);
+  out = fuse(3, out, l3, sizes[1]);
+  out = fuse(2, out, l2, sizes[0]);
+  out = fuse(1, out, l1, [2 * sizes[0][0], 2 * sizes[0][1]]);
+  out = conv(out, 'scratch.output_conv1', features, features / 2, 3, { padding: 1 });
+
+  const up = node('resize', [out], { height, width, mode: 'bilinear', alignCorners: true });
+  const read = conv(
+    embed(up, features / 2, height, width),
+    'scratch.output_conv2.0',
+    features / 2,
+    32,
+    3,
+    { padding: 1 },
+  );
+  return conv(
+    node('relu', [read]),
+    'scratch.output_conv2.2',
+    32,
+    config.outputDim,
+    1,
+    {},
+    true,
+    output,
+  );
+}

package/src/models/hiera.ts

@@ -0,0 +1,231 @@
+/**
+ * Hiera and its feature pyramid, as SAM 2 runs them: an image in, three levels of features out.
+ *
+ * **Four stages of transformer blocks over a shrinking grid.** A 7×7 convolution at stride 4 takes
+ * the image to a quarter; each later stage's first block halves the grid by pooling its queries 2×2
+ * — within each window for attention, and over the whole grid for the residual, which a projection
+ * widens — so attention reads every key and writes a quarter as many queries. Blocks attend within
+ * windows, padded with zeros after the block's norm where the window does not divide the grid, and
+ * the blocks the configuration names attend over the whole grid instead. A first block partitions by
+ * the previous stage's window and unpartitions by half of it.
+ *
+ * **The positional embedding is computed in the graph, from its two weights**: a table resized
+ * bicubically to the grid, plus a window's table tiled across it. Laid out at conversion it would be
+ * 12.6 MB of the file at half precision; computed, it costs two permutes a run.
+ *
+ * **The pyramid runs top-down only into the levels the configuration names**: each stage's output is
+ * projected to the pyramid's width, and a named level adds the level above it doubled by nearest
+ * resizing. The two finest levels are then narrowed by the mask decoder's own projections, as the
+ * upstream's `get_image_features` narrows them once rather than on every prompt. Norms take ε 1e-6.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+import { hieraAttention } from './hieraAttention.ts';
+
+type Four = readonly [number, number, number, number];
+
+export interface HieraConfig {
+  readonly dims: Four;
+  readonly blocks: Four;
+  readonly heads: Four;
+  readonly windows: Four;
+  /** Block indices, counted across stages, that attend over the whole grid. */
+  readonly global: readonly number[];
+  /** The positional table's trained size, square. */
+  readonly background: number;
+  readonly mlpRatio: number;
+  /** The pyramid's width, and the levels (by stage) that take the level above them. */
+  readonly fpn: number;
+  readonly topDown: readonly number[];
+}
+
+export interface HieraOutput {
+  /** `[fpn, grid, grid]` at a sixteenth of the image: the features the decoder and memory read. */
+  readonly features: string;
+  /** `[fpn / 4, 2·grid, 2·grid]` and `[fpn / 8, 4·grid, 4·grid]`: the decoder's skips, narrowed. */
+  readonly high1: string;
+  readonly high0: string;
+}
+
+const EPSILON = 1e-6;
+
+/**
+ * The encoder over `image`, `[3, size, size]` normalised as the upstream prepares it, writing its
+ * three outputs under the names given.
+ */
+export function hiera(
+  weights: Weights,
+  graph: GraphBuilder,
+  config: HieraConfig,
+  image: string,
+  size: number,
+  names: HieraOutput,
+): void {
+  const node = graph.node.bind(graph);
+  const b = 'vision_encoder.backbone.';
+  const read = weights.read.bind(weights);
+  const linear = (x: string, p: string, out: number, inputs: number): string =>
+    node('linear', [x, read(`${p}weight`, [out, inputs]), read(`${p}bias`, [out])]);
+  const norm = (x: string, p: string, dim: number): string =>
+    node('layerNorm', [x, read(`${p}weight`, [dim]), read(`${p}bias`, [dim])], {
+      epsilon: EPSILON,
+    });
+
+  /* The patch embedding, then the positions, both channel-major. */
+  const [d0] = config.dims;
+  let grid = size / 4;
+  const patches = node(
+    'conv2d',
+    [
+      image,
+      read(`${b}patch_embed.projection.weight`, [d0, 3, 7, 7]),
+      read(`${b}patch_embed.projection.bias`, [d0]),
+    ],
+    { stride: 4, padding: 3 },
+  );
+  const window0 = config.windows[0];
+  const tiles = grid / window0;
+  const background = node(
+    'resize',
+    [
+      node('reshape', [read(`${b}pos_embed`, [1, d0, config.background, config.background])], {
+        shape: [d0, config.background, config.background],
+      }),
+    ],
+    { height: grid, width: grid, mode: 'bicubic', alignCorners: false },
+  );
+  /* The window's table added to every tile: tiles to the leading axes, the window to the last. */
+  const byTile = node(
+    'reshape',
+    [
+      node(
+        'permute',
+        [node('reshape', [background], { shape: [d0, tiles, window0, tiles, window0] })],
+        { order: [1, 3, 0, 2, 4] },
+      ),
+    ],
+    { shape: [tiles * tiles, d0 * window0 * window0] },
+  );
+  const window = node('reshape', [read(`${b}pos_embed_window`, [1, d0, window0, window0])], {
+    shape: [d0 * window0 * window0],
+  });
+  const positions = node(
+    'reshape',
+    [
+      node(
+        'permute',
+        [
+          node('reshape', [node('add', [byTile, window])], {
+            shape: [tiles, tiles, d0, window0, window0],
+          }),
+        ],
+        { order: [2, 0, 3, 1, 4] },
+      ),
+    ],
+    { shape: [d0, grid, grid] },
+  );
+  const embedded = node('add', [patches, positions]);
+  /* Tokens, row by row, `[grid · grid, dim]`, from here to the pyramid. */
+  let x = node('permute', [node('reshape', [embedded], { shape: [d0, grid * grid] })], {
+    order: [1, 0],
+  });
+
+  const pooled = (tokens: string, rows: number, dim: number): string => {
+    const map = node('reshape', [node('permute', [tokens], { order: [1, 0] })], {
+      shape: [dim, rows, rows],
+    });
+    const pool = node('maxPool2d', [map], { kernel: 2 });
+    return node('permute', [node('reshape', [pool], { shape: [dim, (rows / 2) * (rows / 2)] })], {
+      order: [1, 0],
+    });
+  };
+
+  const outputs: string[] = [];
+  let index = 0;
+  for (let stage = 0; stage < 4; stage += 1) {
+    for (let block = 0; block < (config.blocks[stage] as number); block += 1) {
+      const p = `${b}blocks.${index}.`;
+      const first = stage > 0 && block === 0;
+      const dim = first ? (config.dims[stage - 1] as number) : (config.dims[stage] as number);
+      const out = config.dims[stage] as number;
+      const heads = config.heads[stage] as number;
+      const side = config.global.includes(index)
+        ? 0
+        : first
+          ? (config.windows[stage - 1] as number)
+          : (config.windows[stage] as number);
+      const normed = norm(x, `${p}layer_norm1.`, dim);
+      const residual = first ? pooled(linear(normed, `${p}proj.`, out, dim), grid, out) : x;
+      const attended = hieraAttention(
+        weights,
+        graph,
+        p,
+        normed,
+        grid,
+        dim,
+        out,
+        heads,
+        side,
+        first,
+      );
+      if (first) grid /= 2;
+      x = node('add', [residual, attended]);
+      const hidden = node('gelu', [
+        linear(norm(x, `${p}layer_norm2.`, out), `${p}mlp.proj_in.`, out * config.mlpRatio, out),
+      ]);
+      x = node('add', [x, linear(hidden, `${p}mlp.proj_out.`, out, out * config.mlpRatio)]);
+      index += 1;
+    }
+    outputs.push(x);
+  }
+
+  /* The pyramid, coarsest first: each stage projected, and a named level adding the one above. */
+  const n = 3;
+  const sides = [size / 4, size / 8, size / 16, size / 32];
+  const levels: string[] = [];
+  let previous = '';
+  for (let i = n; i >= 0; i -= 1) {
+    const p = `vision_encoder.neck.convs.${n - i}.`;
+    const channels = config.dims[i] as number;
+    const side = sides[i] as number;
+    const map = node('reshape', [node('permute', [outputs[i] as string], { order: [1, 0] })], {
+      shape: [channels, side, side],
+    });
+    const down = config.topDown.includes(i) && i !== n;
+    const name = i === 2 ? names.features : undefined;
+    let lateral = node(
+      'conv2d',
+      [map, read(`${p}weight`, [config.fpn, channels, 1, 1]), read(`${p}bias`, [config.fpn])],
+      {},
+      down ? undefined : name,
+    );
+    if (down) {
+      const above = node('resize', [previous], { height: side, width: side, mode: 'nearest' });
+      lateral = node('add', [lateral, above], {}, name);
+    }
+    previous = lateral;
+    levels[i] = lateral;
+  }
+  const d = 'mask_decoder.';
+  node(
+    'conv2d',
+    [
+      levels[0] as string,
+      read(`${d}conv_s0.weight`, [config.fpn / 8, config.fpn, 1, 1]),
+      read(`${d}conv_s0.bias`, [config.fpn / 8]),
+    ],
+    {},
+    names.high0,
+  );
+  node(
+    'conv2d',
+    [
+      levels[1] as string,
+      read(`${d}conv_s1.weight`, [config.fpn / 4, config.fpn, 1, 1]),
+      read(`${d}conv_s1.bias`, [config.fpn / 4]),
+    ],
+    {},
+    names.high1,
+  );
+  /* The coarsest level feeds the one below and is not itself an output, as the upstream keeps three. */
+}

package/src/models/hieraAttention.ts

@@ -0,0 +1,121 @@
+/**
+ * One Hiera block's attention: within windows of its grid, or over the whole grid, its queries
+ * pooled 2×2 within each window where the block changes stage.
+ *
+ * **Windows are partitioned after the block's norm**, padded with zeros where the window does not
+ * divide the grid — so a padded token is zero, where TinyViT's is its norm's β — and unpartitioned
+ * by the queries' window, the padding recomputed for the pooled grid, as the upstream's
+ * `window_unpartition` is given it. The projection lays each head's queries, keys and values out
+ * contiguously, as the attention operator takes them.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+/**
+ * `[rows² or (rows/2)², out]` from `normed`, `[rows², dim]`: the attention under `p` within windows
+ * of `side` — the whole grid for 0 — pooling its queries when `pool`.
+ */
+export function hieraAttention(
+  weights: Weights,
+  graph: GraphBuilder,
+  p: string,
+  normed: string,
+  rows: number,
+  dim: number,
+  out: number,
+  heads: number,
+  side: number,
+  pool: boolean,
+): string {
+  const node = graph.node.bind(graph);
+  const linear = (x: string, q: string, outputs: number, inputs: number): string =>
+    node('linear', [
+      x,
+      weights.read(`${q}weight`, [outputs, inputs]),
+      weights.read(`${q}bias`, [outputs]),
+    ]);
+
+  const span = side === 0 ? rows : side;
+  const pad = (span - (rows % span)) % span;
+  const across = (rows + pad) / span;
+  const count = across * across;
+  let tokens = normed;
+  if (side > 0) {
+    let map = node('reshape', [normed], { shape: [rows, rows, dim] });
+    if (pad > 0) map = node('pad', [map], { after: [pad, pad, 0] });
+    tokens = node(
+      'reshape',
+      [
+        node('permute', [node('reshape', [map], { shape: [across, span, across, span, dim] })], {
+          order: [0, 2, 1, 3, 4],
+        }),
+      ],
+      { shape: [count * span * span, dim] },
+    );
+  }
+  const qkv = linear(tokens, `${p}attn.qkv.`, 3 * out, dim);
+  const part = (i: number): string =>
+    node('slice', [qkv], { axis: 1, start: i * out, end: (i + 1) * out });
+  let queries = node('reshape', [part(0)], { shape: [count, span * span, out] });
+  let qSpan = span;
+  if (pool) {
+    /* Each window's queries as a map, pooled 2×2: windows and channels together as the pool's planes. */
+    const maps = node(
+      'reshape',
+      [
+        node('permute', [node('reshape', [part(0)], { shape: [count, span, span, out] })], {
+          order: [0, 3, 1, 2],
+        }),
+      ],
+      { shape: [count * out, span, span] },
+    );
+    qSpan = span / 2;
+    queries = node(
+      'reshape',
+      [
+        node(
+          'permute',
+          [
+            node('reshape', [node('maxPool2d', [maps], { kernel: 2 })], {
+              shape: [count, out, qSpan, qSpan],
+            }),
+          ],
+          { order: [0, 2, 3, 1] },
+        ),
+      ],
+      { shape: [count, qSpan * qSpan, out] },
+    );
+  }
+  const keys = node('reshape', [part(1)], { shape: [count, span * span, out] });
+  const values = node('reshape', [part(2)], { shape: [count, span * span, out] });
+  const attended = node('attention', [queries, keys, values], { heads });
+  const projected = linear(
+    node('reshape', [attended], { shape: [count * qSpan * qSpan, out] }),
+    `${p}attn.proj.`,
+    out,
+    out,
+  );
+  if (side === 0) return projected;
+  /* Unpartitioned by the queries' window, the padding recomputed for the pooled grid. */
+  const pooledRows = pool ? rows / 2 : rows;
+  const qPad = (qSpan - (pooledRows % qSpan)) % qSpan;
+  const qAcross = (pooledRows + qPad) / qSpan;
+  let map = node(
+    'reshape',
+    [
+      node(
+        'permute',
+        [node('reshape', [projected], { shape: [qAcross, qAcross, qSpan, qSpan, out] })],
+        { order: [0, 2, 1, 3, 4] },
+      ),
+    ],
+    { shape: [pooledRows + qPad, pooledRows + qPad, out] },
+  );
+  if (qPad > 0) {
+    map = node('slice', [node('slice', [map], { axis: 0, start: 0, end: pooledRows })], {
+      axis: 1,
+      start: 0,
+      end: pooledRows,
+    });
+  }
+  return node('reshape', [map], { shape: [pooledRows * pooledRows, out] });
+}