@driftengine/capture 4.0.0

package/src/models/tinyVit.ts

@@ -0,0 +1,178 @@
+/**
+ * TinyViT as MobileSAM runs it: an image in, a 256-channel embedding at a sixteenth of its size out.
+ *
+ * **Four stages, the first convolutional.** A patch embedding of two strided convolutions takes the
+ * image to a quarter; the first stage is inverted residual blocks, each a pointwise expansion, a
+ * depthwise convolution and a pointwise projection; the other three are transformer blocks that
+ * attend within windows, each followed by a depthwise convolution over the grid and an MLP. A merge
+ * between stages widens the channels and halves the grid — except where the upstream keeps its
+ * stride at one, which it decides by the output width and this definition takes as `mergeStrides`.
+ *
+ * **Every convolution carries a batch norm**, folded in at conversion by `convBatchNorm`, and the
+ * transformer blocks attend within windows through `windowAttention`, which pads, partitions and
+ * biases them as the upstream does. The blocks' MLP norms take PyTorch's default ε, 1e-5, and the
+ * neck's channel norms 1e-6.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+import { convBatchNorm } from './batchNorm.ts';
+import { windowAttention } from './windowAttention.ts';
+
+type Four = readonly [number, number, number, number];
+
+export interface TinyVitConfig {
+  readonly embedDims: Four;
+  readonly depths: Four;
+  /** Heads per stage; the first stage is convolutional and has none. */
+  readonly heads: Four;
+  readonly windows: Four;
+  /** The stride of each merge between stages. */
+  readonly mergeStrides: readonly [number, number, number];
+  readonly mlpRatio: number;
+  readonly expandRatio: number;
+  /** The embedding's channels. */
+  readonly neck: number;
+}
+
+const EPSILON = 1e-5;
+const CHANNEL_EPSILON = 1e-6;
+
+/** The grid the encoder's embedding has for an image `size` across. */
+export function tinyVitGrid(config: TinyVitConfig, size: number): number {
+  return config.mergeStrides.reduce((grid, stride) => grid / stride, size / 4);
+}
+
+/**
+ * `[neck, grid, grid]` from `image`, `[3, size, size]`, normalised and padded as the upstream
+ * prepares it; the embedding is written as `output`.
+ */
+export function tinyVit(
+  weights: Weights,
+  graph: GraphBuilder,
+  config: TinyVitConfig,
+  prefix: string,
+  image: string,
+  size: number,
+  output: string,
+): void {
+  const node = graph.node.bind(graph);
+  const dims = config.embedDims;
+  weights.ignore(`${prefix}norm_head.`);
+  weights.ignore(`${prefix}head.`);
+
+  const convBn = (
+    p: string,
+    x: string,
+    cin: number,
+    cout: number,
+    kernel: number,
+    stride: number,
+    groups: number,
+  ): string => convBatchNorm(weights, graph, p, x, cin, cout, kernel, stride, groups);
+  const gelu = (x: string): string => node('gelu', [x]);
+  const tokensOf = (map: string, channels: number, cells: number): string =>
+    node('permute', [node('reshape', [map], { shape: [channels, cells] })], { order: [1, 0] });
+  const mapOf = (tokens: string, channels: number, rows: number, cols: number): string =>
+    node('reshape', [node('permute', [tokens], { order: [1, 0] })], {
+      shape: [channels, rows, cols],
+    });
+  const norm = (x: string, p: string, dim: number, epsilon: number): string =>
+    node('layerNorm', [x, weights.read(`${p}weight`, [dim]), weights.read(`${p}bias`, [dim])], {
+      epsilon,
+    });
+  const linear = (x: string, p: string, out: number, inputs: number): string =>
+    node('linear', [x, weights.read(`${p}weight`, [out, inputs]), weights.read(`${p}bias`, [out])]);
+
+  const merge = (p: string, map: string, dim: number, out: number, stride: number): string => {
+    const expanded = gelu(convBn(`${p}conv1.`, map, dim, out, 1, 1, 1));
+    const strided = gelu(convBn(`${p}conv2.`, expanded, out, out, 3, stride, out));
+    return convBn(`${p}conv3.`, strided, out, out, 1, 1, 1);
+  };
+
+  let map = gelu(convBn(`${prefix}patch_embed.seq.0.`, image, 3, dims[0] / 2, 3, 2, 1));
+  map = convBn(`${prefix}patch_embed.seq.2.`, map, dims[0] / 2, dims[0], 3, 2, 1);
+  let grid = size / 4;
+  const hidden = dims[0] * config.expandRatio;
+  for (let i = 0; i < config.depths[0]; i += 1) {
+    const p = `${prefix}layers.0.blocks.${i}.`;
+    const inner = gelu(convBn(`${p}conv1.`, map, dims[0], hidden, 1, 1, 1));
+    const spread = gelu(convBn(`${p}conv2.`, inner, hidden, hidden, 3, 1, hidden));
+    map = gelu(node('add', [convBn(`${p}conv3.`, spread, hidden, dims[0], 1, 1, 1), map]));
+  }
+  map = merge(`${prefix}layers.0.downsample.`, map, dims[0], dims[1], config.mergeStrides[0]);
+  grid /= config.mergeStrides[0];
+
+  let tokens = tokensOf(map, dims[1], grid * grid);
+  for (let stage = 1; stage < 4; stage += 1) {
+    const dim = dims[stage] as number;
+    const heads = config.heads[stage] as number;
+    const window = config.windows[stage] as number;
+    for (let i = 0; i < (config.depths[stage] as number); i += 1) {
+      const p = `${prefix}layers.${stage}.blocks.${i}.`;
+      const attended = node('add', [
+        tokens,
+        windowAttention(weights, graph, `${p}attn.`, tokens, dim, heads, window, grid, grid),
+      ]);
+      const local = convBn(
+        `${p}local_conv.`,
+        mapOf(attended, dim, grid, grid),
+        dim,
+        dim,
+        3,
+        1,
+        dim,
+      );
+      const mixed = tokensOf(local, dim, grid * grid);
+      const inner = gelu(
+        linear(
+          norm(mixed, `${p}mlp.norm.`, dim, EPSILON),
+          `${p}mlp.fc1.`,
+          dim * config.mlpRatio,
+          dim,
+        ),
+      );
+      tokens = node('add', [mixed, linear(inner, `${p}mlp.fc2.`, dim, dim * config.mlpRatio)]);
+    }
+    if (stage < 3) {
+      const stride = config.mergeStrides[stage] as number;
+      const next = dims[stage + 1] as number;
+      const merged = merge(
+        `${prefix}layers.${stage}.downsample.`,
+        mapOf(tokens, dim, grid, grid),
+        dim,
+        next,
+        stride,
+      );
+      grid /= stride;
+      tokens = tokensOf(merged, next, grid * grid);
+    }
+  }
+
+  /* The neck: a pointwise and a 3×3 convolution, each followed by a norm over the channels. */
+  const channelNorm = (x: string, p: string, name?: string): string =>
+    node(
+      'reshape',
+      [
+        node(
+          'permute',
+          [norm(tokensOf(x, config.neck, grid * grid), p, config.neck, CHANNEL_EPSILON)],
+          { order: [1, 0] },
+        ),
+      ],
+      { shape: [config.neck, grid, grid] },
+      name,
+    );
+  const pointwise = node('conv2d', [
+    mapOf(tokens, dims[3], grid, grid),
+    weights.read(`${prefix}neck.0.weight`, [config.neck, dims[3], 1, 1]),
+  ]);
+  const spatial = node(
+    'conv2d',
+    [
+      channelNorm(pointwise, `${prefix}neck.1.`),
+      weights.read(`${prefix}neck.2.weight`, [config.neck, config.neck, 3, 3]),
+    ],
+    { padding: 1 },
+  );
+  channelNorm(spatial, `${prefix}neck.3.`, output);
+}

package/src/models/twoWay.ts

@@ -0,0 +1,122 @@
+/**
+ * SAM's two-way transformer: a prompt's tokens and an image's, each attending to the other — the
+ * part SAM, MobileSAM and SAM 2 share, under whichever names their checkpoints give it.
+ *
+ * **The tokens attend to themselves, then to the image, then pass an MLP; then the image attends back
+ * to them.** Queries and keys are given their positions afresh at every attention — the tokens as
+ * they were first given, the image by its grid's positions — while values are not. **The first
+ * layer's self-attention replaces the tokens rather than adding to them**, and gives them no
+ * positions, as the upstream's `skip_first_layer_pe` does; every later one is a residual. Attention
+ * between tokens and image narrows to `dim / downsample`. Norms take PyTorch's default ε, 1e-5, and
+ * the MLP's activation is ReLU in every checkpoint here.
+ *
+ * **One definition under two sets of names**, not two definitions: the original's checkpoints say
+ * `out_proj`, `norm1` and `lin1` where Transformers' SAM 2 says `o_proj`, `layer_norm1` and
+ * `proj_in`, and a copy for each would be two implementations of one network, free to drift.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+export interface TwoWayConfig {
+  readonly dim: number;
+  readonly heads: number;
+  readonly mlp: number;
+  readonly depth: number;
+  readonly downsample: number;
+}
+
+/** What the two checkpoint layouts call the transformer's parts. */
+export interface TwoWayNames {
+  /** An attention's output projection. */
+  readonly out: string;
+  /** A layer's `i`-th norm, 1 to 4. */
+  readonly norm: (i: number) => string;
+  readonly finalNorm: string;
+  /** The MLP's two projections. */
+  readonly mlp: readonly [string, string];
+}
+
+/** The original SAM's names, which MobileSAM keeps. */
+export const SAM_NAMES: TwoWayNames = {
+  out: 'out_proj',
+  norm: (i) => `norm${i}`,
+  finalNorm: 'norm_final_attn',
+  mlp: ['lin1', 'lin2'],
+};
+
+/** Transformers' names for SAM 2. */
+export const SAM2_NAMES: TwoWayNames = {
+  out: 'o_proj',
+  norm: (i) => `layer_norm${i}`,
+  finalNorm: 'layer_norm_final_attn',
+  mlp: ['proj_in', 'proj_out'],
+};
+
+const EPSILON = 1e-5;
+
+/**
+ * The transformer under `prefix` over `tokens`, `[count, dim]`, and `image`, `[cells, dim]` in token
+ * layout with `positions` of the same shape; answers both, attended.
+ */
+export function twoWay(
+  weights: Weights,
+  graph: GraphBuilder,
+  prefix: string,
+  config: TwoWayConfig,
+  names: TwoWayNames,
+  tokens: string,
+  image: string,
+  positions: string,
+): { readonly queries: string; readonly keys: string } {
+  const node = graph.node.bind(graph);
+  const { dim, heads } = config;
+  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): string =>
+    node('layerNorm', [x, read(`${p}weight`, [dim]), read(`${p}bias`, [dim])], {
+      epsilon: EPSILON,
+    });
+  const attend = (p: string, q: string, k: string, v: string, narrow: number): string => {
+    const inner = dim / narrow;
+    const attended = node(
+      'attention',
+      [
+        linear(q, `${p}q_proj.`, inner, dim),
+        linear(k, `${p}k_proj.`, inner, dim),
+        linear(v, `${p}v_proj.`, inner, dim),
+      ],
+      { heads },
+    );
+    return linear(attended, `${p}${names.out}.`, dim, inner);
+  };
+  const [into, outOf] = names.mlp;
+
+  let queries = tokens;
+  let keys = image;
+  for (let i = 0; i < config.depth; i += 1) {
+    const p = `${prefix}layers.${i}.`;
+    if (i === 0) {
+      queries = attend(`${p}self_attn.`, queries, queries, queries, 1);
+    } else {
+      const q = node('add', [queries, tokens]);
+      queries = node('add', [queries, attend(`${p}self_attn.`, q, q, queries, 1)]);
+    }
+    queries = norm(queries, `${p}${names.norm(1)}.`);
+    let q = node('add', [queries, tokens]);
+    let k = node('add', [keys, positions]);
+    const toImage = attend(`${p}cross_attn_token_to_image.`, q, k, keys, config.downsample);
+    queries = norm(node('add', [queries, toImage]), `${p}${names.norm(2)}.`);
+    const hidden = node('relu', [linear(queries, `${p}mlp.${into}.`, config.mlp, dim)]);
+    const mixed = linear(hidden, `${p}mlp.${outOf}.`, dim, config.mlp);
+    queries = norm(node('add', [queries, mixed]), `${p}${names.norm(3)}.`);
+    q = node('add', [queries, tokens]);
+    k = node('add', [keys, positions]);
+    const toTokens = attend(`${p}cross_attn_image_to_token.`, k, q, queries, config.downsample);
+    keys = norm(node('add', [keys, toTokens]), `${p}${names.norm(4)}.`);
+  }
+  const q = node('add', [queries, tokens]);
+  const k = node('add', [keys, positions]);
+  const final = attend(`${prefix}final_attn_token_to_image.`, q, k, keys, config.downsample);
+  queries = norm(node('add', [queries, final]), `${prefix}${names.finalNorm}.`);
+  return { queries, keys };
+}

package/src/models/vit.ts

@@ -0,0 +1,265 @@
+/**
+ * Depth Anything 3's backbone: DINOv2 as that model runs it, over one view or several at once.
+ *
+ * **Plain DINOv2 until `altStart`, and something else after it**, which is why this is not a
+ * generic vision transformer. From that block the class token of every view is replaced by a
+ * learned camera token — one for the first view, the reference, and another for the rest — and
+ * blocks alternate: an even one attends within each view, an odd one across all views at once.
+ * From `qkNormStart` each head's queries and keys are layer-normed, and from `ropeStart` they are
+ * turned by a two-dimensional rotary embedding: by the patch's row and column within a view, and
+ * across views by one position shared by every patch, so the global blocks see no geometry but
+ * which view a token belongs to.
+ *
+ * **Each tap is the last within-view output beside the current one**, `2·dim` wide, with only the
+ * second half normed — the upstream's `cat_token` — and its first token kept apart as that view's
+ * camera token, unnormed, which the camera decoder reads.
+ *
+ * **The positional embeddings are resized as the upstream resizes them**: bicubically, by the scale
+ * factor `(grid + 0.1)/trained`, whose inverse is the step a sample moves — not the ratio of the
+ * sizes. Only a square image at the trained grid skips it. The block norms take ε = 1e-6 and every
+ * other norm here PyTorch's default, 1e-5, which is a difference the checkpoint was trained with.
+ *
+ * What it gives up: the upstream reorders views around a chosen reference once there are three or
+ * more; here the first view given is the reference, and choosing it is the caller's.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+import { applyRope, ropeConstants } from './rope.ts';
+
+export interface VitConfig {
+  readonly dim: number;
+  readonly heads: number;
+  readonly depth: number;
+  readonly hidden: number;
+  readonly patch: number;
+  /** The patch grid the positional embeddings were trained at: 518 / 14 = 37. */
+  readonly trainedGrid: number;
+  readonly taps: readonly number[];
+  readonly altStart: number;
+  readonly qkNormStart: number;
+  readonly ropeStart: number;
+}
+
+export interface VitOutput {
+  /** Per tap, per view: the patch tokens, `[rows · cols, 2 · dim]`. */
+  readonly features: readonly (readonly string[])[];
+  /** Per tap: every view's camera token, `[views, 2 · dim]`. */
+  readonly cameras: readonly string[];
+}
+
+/** Block norms are trained with ε = 1e-6; the rest keep PyTorch's default. */
+const BLOCK_EPSILON = 1e-6;
+const DEFAULT_EPSILON = 1e-5;
+
+export function vitBackbone(
+  weights: Weights,
+  graph: GraphBuilder,
+  config: VitConfig,
+  prefix: string,
+  images: readonly string[],
+  rows: number,
+  cols: number,
+): VitOutput {
+  const { dim, heads } = config;
+  const headDim = dim / heads;
+  const patches = rows * cols;
+  const tokens = patches + 1;
+  const w = (name: string, shape: readonly number[]): string =>
+    weights.read(`${prefix}${name}`, shape);
+  const node = graph.node.bind(graph);
+
+  /* Positions, resized to this grid unless it is the trained one on a square image. */
+  const trained = config.trainedGrid;
+  const table = node('reshape', [w('pos_embed', [1, trained * trained + 1, dim])], {
+    shape: [trained * trained + 1, dim],
+  });
+  let positions = table;
+  if (rows !== trained || cols !== trained) {
+    const grid = node(
+      'permute',
+      [
+        node(
+          'reshape',
+          [node('slice', [table], { axis: 0, start: 1, end: trained * trained + 1 })],
+          {
+            shape: [trained, trained, dim],
+          },
+        ),
+      ],
+      { order: [2, 0, 1] },
+    );
+    const resized = node('resize', [grid], {
+      height: rows,
+      width: cols,
+      mode: 'bicubic',
+      alignCorners: false,
+      stepHeight: trained / (rows + 0.1),
+      stepWidth: trained / (cols + 0.1),
+    });
+    const flat = node('reshape', [node('permute', [resized], { order: [1, 2, 0] })], {
+      shape: [patches, dim],
+    });
+    positions = node('concat', [node('slice', [table], { axis: 0, start: 0, end: 1 }), flat], {
+      axis: 0,
+    });
+  }
+
+  const cls = node('reshape', [w('cls_token', [1, 1, dim])], { shape: [1, dim] });
+  let views = images.map((image) => {
+    const embedded = node(
+      'patchEmbed',
+      [
+        image,
+        w('patch_embed.proj.weight', [dim, 3, config.patch, config.patch]),
+        w('patch_embed.proj.bias', [dim]),
+      ],
+      { patch: config.patch },
+    );
+    return node('add', [node('concat', [cls, embedded], { axis: 0 }), positions]);
+  });
+
+  /* Rotary tables: within a view by row and column from 1, across views all patches at 1, 1. */
+  const local = new Int32Array(2 * tokens);
+  const global = new Int32Array(2 * tokens * images.length);
+  for (let p = 0; p < patches; p += 1) {
+    local[2 * (p + 1)] = Math.floor(p / cols) + 1;
+    local[2 * (p + 1) + 1] = (p % cols) + 1;
+  }
+  for (let v = 0; v < images.length; v += 1) {
+    for (let p = 1; p < tokens; p += 1)
+      global.fill(1, 2 * (v * tokens + p), 2 * (v * tokens + p) + 2);
+  }
+  const rope =
+    config.ropeStart >= 0
+      ? {
+          local: ropeConstants(weights, 'rope.local', local, heads, headDim),
+          global: ropeConstants(weights, 'rope.global', global, heads, headDim),
+        }
+      : null;
+
+  const block = (i: number, x: string, count: number, across: boolean): string => {
+    const p = `blocks.${i}.`;
+    const h = node('layerNorm', [x, w(`${p}norm1.weight`, [dim]), w(`${p}norm1.bias`, [dim])], {
+      epsilon: BLOCK_EPSILON,
+    });
+    const qkv = node('linear', [
+      h,
+      w(`${p}attn.qkv.weight`, [3 * dim, dim]),
+      w(`${p}attn.qkv.bias`, [3 * dim]),
+    ]);
+    let q = node('slice', [qkv], { axis: 1, start: 0, end: dim });
+    let k = node('slice', [qkv], { axis: 1, start: dim, end: 2 * dim });
+    const v = node('slice', [qkv], { axis: 1, start: 2 * dim, end: 3 * dim });
+    if (config.qkNormStart >= 0 && i >= config.qkNormStart) {
+      const perHead = (value: string, norm: string): string =>
+        node(
+          'reshape',
+          [
+            node(
+              'layerNorm',
+              [
+                node('reshape', [value], { shape: [count * heads, headDim] }),
+                w(`${p}attn.${norm}.weight`, [headDim]),
+                w(`${p}attn.${norm}.bias`, [headDim]),
+              ],
+              { epsilon: DEFAULT_EPSILON },
+            ),
+          ],
+          { shape: [count, dim] },
+        );
+      q = perHead(q, 'q_norm');
+      k = perHead(k, 'k_norm');
+    }
+    if (rope !== null && i >= config.ropeStart) {
+      const tables = across ? rope.global : rope.local;
+      q = applyRope(graph, q, count, heads, headDim, tables.cos, tables.sin);
+      k = applyRope(graph, k, count, heads, headDim, tables.cos, tables.sin);
+    }
+    const attended = node('attention', [q, k, v], { heads });
+    const projected = node('linear', [
+      attended,
+      w(`${p}attn.proj.weight`, [dim, dim]),
+      w(`${p}attn.proj.bias`, [dim]),
+    ]);
+    const x1 = node('add', [x, node('mul', [projected, w(`${p}ls1.gamma`, [dim])])]);
+    const h2 = node('layerNorm', [x1, w(`${p}norm2.weight`, [dim]), w(`${p}norm2.bias`, [dim])], {
+      epsilon: BLOCK_EPSILON,
+    });
+    const inner = node('gelu', [
+      node('linear', [
+        h2,
+        w(`${p}mlp.fc1.weight`, [config.hidden, dim]),
+        w(`${p}mlp.fc1.bias`, [config.hidden]),
+      ]),
+    ]);
+    const out = node('linear', [
+      inner,
+      w(`${p}mlp.fc2.weight`, [dim, config.hidden]),
+      w(`${p}mlp.fc2.bias`, [dim]),
+    ]);
+    return node('add', [x1, node('mul', [out, w(`${p}ls2.gamma`, [dim])])]);
+  };
+
+  const features: string[][] = [];
+  const cameras: string[] = [];
+  let latest = views;
+  for (let i = 0; i < config.depth; i += 1) {
+    if (config.altStart >= 0 && i === config.altStart) {
+      const camera = node('reshape', [w('camera_token', [1, 2, dim])], { shape: [2, dim] });
+      const reference = node('slice', [camera], { axis: 0, start: 0, end: 1 });
+      const source = node('slice', [camera], { axis: 0, start: 1, end: 2 });
+      views = views.map((x, v) =>
+        node(
+          'concat',
+          [v === 0 ? reference : source, node('slice', [x], { axis: 0, start: 1, end: tokens })],
+          { axis: 0 },
+        ),
+      );
+    }
+    const across = config.altStart >= 0 && i >= config.altStart && i % 2 === 1;
+    if (across) {
+      const all = views.length === 1 ? (views[0] as string) : node('concat', views, { axis: 0 });
+      const mixed = block(i, all, tokens * views.length, true);
+      views = views.map((_, v) =>
+        views.length === 1
+          ? mixed
+          : node('slice', [mixed], { axis: 0, start: v * tokens, end: (v + 1) * tokens }),
+      );
+    } else {
+      views = views.map((x) => block(i, x, tokens, false));
+      latest = views;
+    }
+    if (config.taps.includes(i)) {
+      const tapped = views.map((x, v) => node('concat', [latest[v] as string, x], { axis: 1 }));
+      cameras.push(
+        tapped.length === 1
+          ? node('slice', [tapped[0] as string], { axis: 0, start: 0, end: 1 })
+          : node(
+              'concat',
+              tapped.map((t) => node('slice', [t], { axis: 0, start: 0, end: 1 })),
+              { axis: 0 },
+            ),
+      );
+      features.push(
+        tapped.map((t) => {
+          const normed = node(
+            'layerNorm',
+            [
+              node('slice', [t], { axis: 1, start: dim, end: 2 * dim }),
+              w('norm.weight', [dim]),
+              w('norm.bias', [dim]),
+            ],
+            { epsilon: DEFAULT_EPSILON },
+          );
+          const joined = node(
+            'concat',
+            [node('slice', [t], { axis: 1, start: 0, end: dim }), normed],
+            { axis: 1 },
+          );
+          return node('slice', [joined], { axis: 0, start: 1, end: tokens });
+        }),
+      );
+    }
+  }
+  return { features, cameras };
+}