@driftengine/capture 4.0.0

package/src/models/clipTower.ts

@@ -0,0 +1,87 @@
+/**
+ * CLIP's encoder tower, which OWLv2's image and text graphs are each one of.
+ *
+ * **Pre-norm layers**: a layer norm, attention over the row's own window, and the projection added
+ * back; then a layer norm, a perceptron, and that added back. The activation is the upstream's
+ * quick GELU, `x·σ(1.702·x)`, which is the logistic and two multiplies rather than the error
+ * function `gelu` computes.
+ *
+ * **A causal tower masks its scores** with −10⁴ above the diagonal, where the upstream uses the
+ * least finite single. Both give a masked key a weight of exactly zero — exp of anything below
+ * −104 is zero in single precision, and scores stay within a few hundred of each other — and −10⁴
+ * is exact at half precision, where the upstream's value would be infinite and a device may assume
+ * no infinities. **What would make it wrong** is scores reaching thousands, which would need a
+ * checkpoint whose norms let them.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+export interface ClipTowerConfig {
+  readonly dim: number;
+  readonly heads: number;
+  readonly depth: number;
+  readonly hidden: number;
+}
+
+const EPSILON = 1e-5;
+const MASKED = -1e4;
+
+/**
+ * CLIP's layers over `x`, `[batch · rows, dim]`, each batch's rows attending among themselves —
+ * causally when `causal` is set.
+ */
+export function clipLayers(
+  weights: Weights,
+  graph: GraphBuilder,
+  prefix: string,
+  config: ClipTowerConfig,
+  x: string,
+  batch: number,
+  rows: number,
+  causal: boolean,
+): string {
+  const { dim, heads, hidden } = config;
+  const node = graph.node.bind(graph);
+  const w = (name: string, shape: readonly number[]): string =>
+    weights.read(`${prefix}${name}`, shape);
+  const linear = (value: string, p: string, out: number, inputs: number): string =>
+    node('linear', [value, w(`${p}.weight`, [out, inputs]), w(`${p}.bias`, [out])]);
+  const norm = (value: string, p: string): string =>
+    node('layerNorm', [value, w(`${p}.weight`, [dim]), w(`${p}.bias`, [dim])], {
+      epsilon: EPSILON,
+    });
+  const quick = weights.constant(
+    `quick_gelu.${hidden}`,
+    [hidden],
+    new Float32Array(hidden).fill(1.702),
+  );
+  let mask: string | null = null;
+  if (causal) {
+    const table = new Float32Array(heads * rows * rows);
+    for (let h = 0; h < heads; h += 1) {
+      for (let q = 0; q < rows; q += 1) {
+        for (let k = q + 1; k < rows; k += 1) table[(h * rows + q) * rows + k] = MASKED;
+      }
+    }
+    mask = weights.constant(`causal.${heads}x${rows}`, [heads, rows, rows], table);
+  }
+  const batched = (value: string): string =>
+    batch > 1 || mask !== null ? node('reshape', [value], { shape: [batch, rows, dim] }) : value;
+
+  let out = x;
+  for (let i = 0; i < config.depth; i += 1) {
+    const p = `encoder.layers.${i}.`;
+    const h = norm(out, `${p}layer_norm1`);
+    const q = batched(linear(h, `${p}self_attn.q_proj`, dim, dim));
+    const k = batched(linear(h, `${p}self_attn.k_proj`, dim, dim));
+    const v = batched(linear(h, `${p}self_attn.v_proj`, dim, dim));
+    let attended = node('attention', mask === null ? [q, k, v] : [q, k, v, mask], { heads });
+    if (batch > 1 || mask !== null) {
+      attended = node('reshape', [attended], { shape: [batch * rows, dim] });
+    }
+    out = node('add', [out, linear(attended, `${p}self_attn.out_proj`, dim, dim)]);
+    const inner = linear(norm(out, `${p}layer_norm2`), `${p}mlp.fc1`, hidden, dim);
+    const activated = node('mul', [inner, node('sigmoid', [node('mul', [inner, quick])])]);
+    out = node('add', [out, linear(activated, `${p}mlp.fc2`, dim, hidden)]);
+  }
+  return out;
+}

package/src/models/depthAnything2.ts

@@ -0,0 +1,276 @@
+/**
+ * Depth Anything V2 Small, the lighter fallback: plain DINOv2 and a DPT, in the Transformers layout
+ * its Apache-2.0 checkpoint is published in.
+ *
+ * **Plain where Depth Anything 3 is not**: one view, no camera token, no query-key norms, no rotary
+ * embedding, and positions resized to the grid by the ratio of the sizes, as Transformers' DINOv2
+ * resizes them. Four taps — after the third, sixth, ninth and twelfth blocks, each through the
+ * backbone's last norm — are reassembled into a pyramid and fused as Depth Anything 3's head fuses
+ * its own, without the sinusoidal positions that head adds. The output is one channel through a
+ * rectifier: relative inverse depth, larger nearer, with no scale.
+ *
+ * **The checkpoint keeps Transformers' older names** — `attention.attention.query`,
+ * `layer_scale1.lambda1` — which the library renames as it loads; this reads them as stored. Its
+ * mask token is for training and set aside, and so is the first fusion block's lateral unit, which
+ * the upstream builds and never runs because the coarsest level has no level below it.
+ */
+import type { Architecture } from '@driftengine/texture';
+
+export interface DepthAnything2Config {
+  readonly dim: number;
+  readonly heads: number;
+  readonly depth: number;
+  readonly hidden: number;
+  readonly patch: number;
+  readonly trainedGrid: number;
+  /** The blocks whose outputs are tapped, counted from zero. */
+  readonly taps: readonly [number, number, number, number];
+  readonly neck: readonly [number, number, number, number];
+  readonly fusion: number;
+  readonly headHidden: number;
+}
+
+/** The checkpoint's configuration at its pinned revision. */
+export const DEPTH_ANYTHING_2_SMALL: DepthAnything2Config = {
+  dim: 384,
+  heads: 6,
+  depth: 12,
+  hidden: 1536,
+  patch: 14,
+  trainedGrid: 37,
+  taps: [2, 5, 8, 11],
+  neck: [48, 96, 192, 384],
+  fusion: 64,
+  headHidden: 32,
+};
+
+/** Transformers' DINOv2 trains every norm at ε = 1e-6. */
+const EPSILON = 1e-6;
+
+/**
+ * The graph for one image of `height × width` pixels, each a multiple of the patch, normalised by
+ * ImageNet's mean and deviation: input `image`, `[3, height, width]`; output `depth`,
+ * `[1, height, width]`.
+ */
+export function depthAnything2(
+  config: DepthAnything2Config,
+  height: number,
+  width: number,
+): Architecture {
+  const rows = height / config.patch;
+  const cols = width / config.patch;
+  if (!Number.isInteger(rows) || !Number.isInteger(cols)) {
+    throw new RangeError(
+      `${height} × ${width} is not a whole number of ${config.patch}-pixel patches`,
+    );
+  }
+  return (weights, graph) => {
+    const { dim, heads, trainedGrid } = config;
+    const node = graph.node.bind(graph);
+    const w = (name: string, shape: readonly number[]): string => weights.read(name, shape);
+    const linear = (x: string, name: string, to: number, from: number): string =>
+      node('linear', [x, w(`${name}.weight`, [to, from]), w(`${name}.bias`, [to])]);
+    const conv = (
+      x: string,
+      name: string,
+      to: number,
+      from: number,
+      kernel: number,
+      attributes = {},
+      bias = true,
+    ): 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,
+      );
+    weights.ignore('backbone.embeddings.mask_token');
+    weights.ignore('neck.fusion_stage.layers.0.residual_layer1.');
+
+    const e = 'backbone.embeddings.';
+    const patches = rows * cols;
+    const tokens = patches + 1;
+    const table = node(
+      'reshape',
+      [w(`${e}position_embeddings`, [1, trainedGrid * trainedGrid + 1, dim])],
+      {
+        shape: [trainedGrid * trainedGrid + 1, dim],
+      },
+    );
+    let positions = table;
+    if (rows !== trainedGrid || cols !== trainedGrid) {
+      const grid = node(
+        'permute',
+        [
+          node(
+            'reshape',
+            [node('slice', [table], { axis: 0, start: 1, end: trainedGrid * trainedGrid + 1 })],
+            {
+              shape: [trainedGrid, trainedGrid, dim],
+            },
+          ),
+        ],
+        { order: [2, 0, 1] },
+      );
+      const resized = node('resize', [grid], {
+        height: rows,
+        width: cols,
+        mode: 'bicubic',
+        alignCorners: false,
+      });
+      positions = node(
+        'concat',
+        [
+          node('slice', [table], { axis: 0, start: 0, end: 1 }),
+          node('reshape', [node('permute', [resized], { order: [1, 2, 0] })], {
+            shape: [patches, dim],
+          }),
+        ],
+        { axis: 0 },
+      );
+    }
+    const embedded = node(
+      'patchEmbed',
+      [
+        'image',
+        w(`${e}patch_embeddings.projection.weight`, [dim, 3, config.patch, config.patch]),
+        w(`${e}patch_embeddings.projection.bias`, [dim]),
+      ],
+      { patch: config.patch },
+    );
+    let x = node('add', [
+      node(
+        'concat',
+        [node('reshape', [w(`${e}cls_token`, [1, 1, dim])], { shape: [1, dim] }), embedded],
+        { axis: 0 },
+      ),
+      positions,
+    ]);
+
+    const taps: string[] = [];
+    for (let i = 0; i < config.depth; i += 1) {
+      const p = `backbone.encoder.layer.${i}.`;
+      const h = node('layerNorm', [x, w(`${p}norm1.weight`, [dim]), w(`${p}norm1.bias`, [dim])], {
+        epsilon: EPSILON,
+      });
+      const attended = node(
+        'attention',
+        [
+          linear(h, `${p}attention.attention.query`, dim, dim),
+          linear(h, `${p}attention.attention.key`, dim, dim),
+          linear(h, `${p}attention.attention.value`, dim, dim),
+        ],
+        { heads },
+      );
+      const out = linear(attended, `${p}attention.output.dense`, dim, dim);
+      x = node('add', [x, node('mul', [out, w(`${p}layer_scale1.lambda1`, [dim])])]);
+      const h2 = node('layerNorm', [x, w(`${p}norm2.weight`, [dim]), w(`${p}norm2.bias`, [dim])], {
+        epsilon: EPSILON,
+      });
+      const mlp = linear(
+        node('gelu', [linear(h2, `${p}mlp.fc1`, config.hidden, dim)]),
+        `${p}mlp.fc2`,
+        dim,
+        config.hidden,
+      );
+      x = node('add', [x, node('mul', [mlp, w(`${p}layer_scale2.lambda1`, [dim])])]);
+      if (config.taps.includes(i)) {
+        taps.push(
+          node(
+            'layerNorm',
+            [x, w('backbone.layernorm.weight', [dim]), w('backbone.layernorm.bias', [dim])],
+            { epsilon: EPSILON },
+          ),
+        );
+      }
+    }
+
+    const half = (size: number): number => Math.floor((size - 1) / 2) + 1;
+    const sizes: [number, number][] = [
+      [4 * rows, 4 * cols],
+      [2 * rows, 2 * cols],
+      [rows, cols],
+      [half(rows), half(cols)],
+    ];
+    const levels = taps.map((tap, k) => {
+      const channels = config.neck[k] as number;
+      const r = `neck.reassemble_stage.layers.${k}.`;
+      const image = node(
+        'reshape',
+        [
+          node('permute', [node('slice', [tap], { axis: 0, start: 1, end: tokens })], {
+            order: [1, 0],
+          }),
+        ],
+        {
+          shape: [dim, rows, cols],
+        },
+      );
+      const projected = conv(image, `${r}projection`, channels, dim, 1);
+      const resized =
+        k === 0
+          ? node(
+              'convTranspose2d',
+              [
+                projected,
+                w(`${r}resize.weight`, [channels, channels, 4, 4]),
+                w(`${r}resize.bias`, [channels]),
+              ],
+              { stride: 4 },
+            )
+          : k === 1
+            ? node(
+                'convTranspose2d',
+                [
+                  projected,
+                  w(`${r}resize.weight`, [channels, channels, 2, 2]),
+                  w(`${r}resize.bias`, [channels]),
+                ],
+                { stride: 2 },
+              )
+            : k === 2
+              ? projected
+              : conv(projected, `${r}resize`, channels, channels, 3, { stride: 2, padding: 1 });
+      return conv(resized, `neck.convs.${k}`, config.fusion, channels, 3, { padding: 1 }, false);
+    });
+
+    const f = config.fusion;
+    const residual = (value: string, name: string): string => {
+      const first = conv(node('relu', [value]), `${name}.convolution1`, f, f, 3, { padding: 1 });
+      return node('add', [
+        conv(node('relu', [first]), `${name}.convolution2`, f, f, 3, { padding: 1 }),
+        value,
+      ]);
+    };
+    /* Coarsest first: layer 0 takes the fourth level alone, and each after it adds the next finer. */
+    let fused = '';
+    for (let n = 0; n < 4; n += 1) {
+      const level = levels[3 - n] as string;
+      const name = `neck.fusion_stage.layers.${n}`;
+      const merged =
+        n === 0 ? level : node('add', [fused, residual(level, `${name}.residual_layer1`)]);
+      const refined = residual(merged, `${name}.residual_layer2`);
+      const size =
+        n < 3
+          ? (sizes[2 - n] as [number, number])
+          : ([2 * sizes[0][0], 2 * sizes[0][1]] as [number, number]);
+      const up = node('resize', [refined], {
+        height: size[0],
+        width: size[1],
+        mode: 'bilinear',
+        alignCorners: true,
+      });
+      fused = conv(up, `${name}.projection`, f, f, 1);
+    }
+
+    const first = conv(fused, 'head.conv1', f / 2, f, 3, { padding: 1 });
+    const up = node('resize', [first], { height, width, mode: 'bilinear', alignCorners: true });
+    const second = node('relu', [
+      conv(up, 'head.conv2', config.headHidden, f / 2, 3, { padding: 1 }),
+    ]);
+    node('relu', [conv(second, 'head.conv3', 1, config.headHidden, 1)], {}, 'depth');
+    return { inputs: [{ name: 'image', shape: [3, height, width] }], outputs: ['depth'] };
+  };
+}

package/src/models/depthAnything3.ts

@@ -0,0 +1,229 @@
+/**
+ * Depth Anything 3, as a definition the runtime builds a graph from at whatever size a clip needs.
+ *
+ * **From one to a handful of views at once, a depth map, its confidence and a camera per view.**
+ * The backbone is `vit.ts`, the head `dualDpt.ts`, and the camera decoder here: two layers over
+ * each view's camera token from the last tap, then a translation, a scalar-last quaternion and two
+ * fields of view. The graph ends at logits and at that pose encoding; `decodeDepth` and
+ * `decodeCamera` turn them into metres of relative depth, a confidence of at least one, a
+ * world-to-camera transform and a pinhole, as the upstream's own code does after its network.
+ *
+ * **What is set aside, by name**: the camera encoder, which conditions on poses a caller already
+ * has and this capture never does; and the head's ray branch, whose output the upstream discards
+ * whenever the camera decoder is present. Converting a checkpoint that holds anything else fails
+ * naming it.
+ *
+ * **What would make this wrong**: a checkpoint trained with another block layout — the upstream's
+ * large and giant models use a different feed-forward — which their configurations would say, and
+ * which are refused for their licences in any case.
+ */
+import { exactCos, exactExp, exactSin } from '@driftengine/core';
+import type { Architecture } from '@driftengine/texture';
+
+import { dualDptHead, type HeadConfig } from './dualDpt.ts';
+import { vitBackbone, type VitConfig } from './vit.ts';
+
+export interface DepthAnything3Config {
+  readonly backbone: VitConfig;
+  readonly head: HeadConfig;
+}
+
+const PATCH = 14;
+
+/** The two accepted sizes, from each checkpoint's own configuration at its pinned revision. */
+export const DEPTH_ANYTHING_3: Readonly<Record<'small' | 'base', DepthAnything3Config>> = {
+  small: {
+    backbone: {
+      dim: 384,
+      heads: 6,
+      depth: 12,
+      hidden: 1536,
+      patch: PATCH,
+      trainedGrid: 37,
+      taps: [5, 7, 9, 11],
+      altStart: 4,
+      qkNormStart: 4,
+      ropeStart: 4,
+    },
+    head: { dimIn: 768, features: 64, outChannels: [48, 96, 192, 384], outputDim: 2, patch: PATCH },
+  },
+  base: {
+    backbone: {
+      dim: 768,
+      heads: 12,
+      depth: 12,
+      hidden: 3072,
+      patch: PATCH,
+      trainedGrid: 37,
+      taps: [5, 7, 9, 11],
+      altStart: 4,
+      qkNormStart: 4,
+      ropeStart: 4,
+    },
+    head: {
+      dimIn: 1536,
+      features: 128,
+      outChannels: [96, 192, 384, 768],
+      outputDim: 2,
+      patch: PATCH,
+    },
+  },
+};
+
+/** What the upstream checkpoints hold and no graph here reads. */
+const SET_ASIDE = [
+  'model.cam_enc.',
+  'model.head.scratch.output_conv1_aux.',
+  'model.head.scratch.output_conv2_aux.',
+  'model.head.scratch.refinenet1_aux.',
+  'model.head.scratch.refinenet2_aux.',
+  'model.head.scratch.refinenet3_aux.',
+  'model.head.scratch.refinenet4_aux.',
+];
+
+/**
+ * The graph for `views` images of `height × width` pixels, each a multiple of 14, normalised by
+ * ImageNet's mean and deviation. Inputs are `image0` … as `[3, height, width]`; outputs are
+ * `logits0` … as `[2, height, width]` and `pose` as `[views, 9]`. The first view is the reference.
+ */
+export function depthAnything3(
+  config: DepthAnything3Config,
+  views: number,
+  height: number,
+  width: number,
+): Architecture {
+  const rows = height / PATCH;
+  const cols = width / PATCH;
+  if (!Number.isInteger(rows) || !Number.isInteger(cols)) {
+    throw new RangeError(`${height} × ${width} is not a whole number of ${PATCH}-pixel patches`);
+  }
+  return (weights, graph) => {
+    for (const prefix of SET_ASIDE) weights.ignore(prefix);
+    const images = Array.from({ length: views }, (_, v) => `image${v}`);
+    const { features, cameras } = vitBackbone(
+      weights,
+      graph,
+      config.backbone,
+      'model.backbone.pretrained.',
+      images,
+      rows,
+      cols,
+    );
+    const outputs = images.map((_, v) =>
+      dualDptHead(
+        weights,
+        graph,
+        config.head,
+        'model.head.',
+        features.map((tap) => tap[v] as string),
+        rows,
+        cols,
+        `logits${v}`,
+      ),
+    );
+
+    const width2 = 2 * config.backbone.dim;
+    const node = graph.node.bind(graph);
+    const layer = (x: string, name: string, to: number): string =>
+      node('linear', [
+        x,
+        weights.read(`model.cam_dec.${name}.weight`, [to, width2]),
+        weights.read(`model.cam_dec.${name}.bias`, [to]),
+      ]);
+    const trunk = node('relu', [
+      layer(
+        node('relu', [layer(cameras[cameras.length - 1] as string, 'backbone.0', width2)]),
+        'backbone.2',
+        width2,
+      ),
+    ]);
+    node(
+      'concat',
+      [
+        layer(trunk, 'fc_t', 3),
+        layer(trunk, 'fc_qvec', 4),
+        node('relu', [layer(trunk, 'fc_fov.0', 2)]),
+      ],
+      { axis: 1 },
+      'pose',
+    );
+    return {
+      inputs: images.map((name) => ({ name, shape: [3, height, width] })),
+      outputs: [...outputs, 'pose'],
+    };
+  };
+}
+
+/**
+ * Depth and confidence from one view's logits, `[2, height, width]`: `exp` of the first channel,
+ * and one more than `exp` of the second, into the caller's maps.
+ */
+export function decodeDepth(
+  logits: Float32Array,
+  depth: Float32Array,
+  confidence: Float32Array,
+): void {
+  const pixels = depth.length;
+  for (let i = 0; i < pixels; i += 1) {
+    depth[i] = exactExp(logits[i] as number);
+    confidence[i] = exactExp(logits[pixels + i] as number) + 1;
+  }
+}
+
+/**
+ * One view's camera from its pose encoding: `worldToCamera`, 3 × 4 row-major — the inverse of the
+ * camera-to-world rotation and translation the decoder predicts — and `intrinsics`, 3 × 3 row-major,
+ * with focal lengths from the two fields of view and the principal point at the centre.
+ */
+export function decodeCamera(
+  pose: Float32Array,
+  at: number,
+  height: number,
+  width: number,
+  worldToCamera: Float32Array,
+  intrinsics: Float32Array,
+): void {
+  const p = at * 9;
+  const tx = pose[p] as number;
+  const ty = pose[p + 1] as number;
+  const tz = pose[p + 2] as number;
+  const i = pose[p + 3] as number;
+  const j = pose[p + 4] as number;
+  const k = pose[p + 5] as number;
+  const r = pose[p + 6] as number;
+  const s = 2 / (i * i + j * j + k * k + r * r);
+  /*
+   * The camera-to-world rotation is `quat_to_mat`'s; its inverse is the transpose, so row `n` here
+   * is column `n` there, and the translation is carried back through it.
+   */
+  const r00 = 1 - s * (j * j + k * k);
+  const r01 = s * (i * j - k * r);
+  const r02 = s * (i * k + j * r);
+  const r10 = s * (i * j + k * r);
+  const r11 = 1 - s * (i * i + k * k);
+  const r12 = s * (j * k - i * r);
+  const r20 = s * (i * k - j * r);
+  const r21 = s * (j * k + i * r);
+  const r22 = 1 - s * (i * i + j * j);
+  worldToCamera[0] = r00;
+  worldToCamera[1] = r10;
+  worldToCamera[2] = r20;
+  worldToCamera[3] = -(r00 * tx + r10 * ty + r20 * tz);
+  worldToCamera[4] = r01;
+  worldToCamera[5] = r11;
+  worldToCamera[6] = r21;
+  worldToCamera[7] = -(r01 * tx + r11 * ty + r21 * tz);
+  worldToCamera[8] = r02;
+  worldToCamera[9] = r12;
+  worldToCamera[10] = r22;
+  worldToCamera[11] = -(r02 * tx + r12 * ty + r22 * tz);
+  const tangent = (angle: number): number => exactSin(angle / 2) / exactCos(angle / 2);
+  const fy = height / 2 / Math.max(tangent(pose[p + 7] as number), 1e-6);
+  const fx = width / 2 / Math.max(tangent(pose[p + 8] as number), 1e-6);
+  intrinsics.fill(0);
+  intrinsics[0] = fx;
+  intrinsics[2] = width / 2;
+  intrinsics[4] = fy;
+  intrinsics[5] = height / 2;
+  intrinsics[8] = 1;
+}