@driftengine/capture 4.0.0

package/src/models/samDecoder.ts

@@ -0,0 +1,225 @@
+/**
+ * SAM's mask decoder: an image embedding and a prompt's tokens in, masks and their predicted
+ * quality out — MobileSAM's, and the original SAM's, which share it.
+ *
+ * **A two-way transformer** (`twoWay.ts`, under the original's names) over the prompt's tokens,
+ * behind a learned token for the quality estimate and one per mask, and the image with its dense
+ * prompt added.
+ *
+ * **Masks come from a product, not a convolution**: the image is upsampled four times by two
+ * transposed convolutions, and each mask token, through an MLP of its own, becomes a vector that
+ * is dotted with every upsampled pixel. The first mask is the one to take for an unambiguous
+ * prompt, and the other three for an ambiguous one, as the upstream's `multimask_output` chooses.
+ *
+ * **Every prompt weight is read, whether this graph's nodes use it or not**: the prompt's tokens
+ * are made on the host from the Fourier weights, and the mask input's downscaling is only in the
+ * graph that takes one — but the converted file is built from one graph, and each has to find
+ * every weight the others need in it.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+import { samGridPositions } from './samPrompt.ts';
+import { SAM_NAMES, twoWay } from './twoWay.ts';
+
+export interface SamDecoderConfig {
+  readonly dim: number;
+  readonly heads: number;
+  readonly mlp: number;
+  readonly depth: number;
+  /** The width attention between tokens and image narrows by. */
+  readonly downsample: number;
+  /** The mask tokens: one for an unambiguous prompt and three for an ambiguous one. */
+  readonly masks: number;
+  readonly iouHidden: number;
+  readonly iouDepth: number;
+  /** The channels a mask input is downscaled through. */
+  readonly maskChannels: number;
+}
+
+const CHANNEL_EPSILON = 1e-6;
+
+/** What the decoder's two outputs are named. */
+export interface SamDecoderOutput {
+  /** `[masks, 4·grid, 4·grid]`: each mask's logits at four times the embedding's grid. */
+  readonly masks: string;
+  /** `[1, masks]`: each mask's predicted intersection over union. */
+  readonly quality: string;
+}
+
+/**
+ * The decoder over `embedding`, `[dim, grid, grid]`, and `promptTokens`, `[tokens, dim]` or null for a
+ * prompt of none, with `mask`, `[1, 4·grid, 4·grid]`, the previous logits to refine, or null.
+ */
+export function samDecoder(
+  weights: Weights,
+  graph: GraphBuilder,
+  config: SamDecoderConfig,
+  embedding: string,
+  promptTokens: string | null,
+  mask: string | null,
+  grid: number,
+  names: SamDecoderOutput,
+): void {
+  const node = graph.node.bind(graph);
+  const { dim, masks } = config;
+  const cells = grid * grid;
+  const d = 'mask_decoder.';
+  const t = `${d}transformer.`;
+  const pe = 'prompt_encoder.';
+  const read = weights.read.bind(weights);
+
+  const linear = (x: string, p: string, out: number, inputs: number, name?: string): string =>
+    node('linear', [x, read(`${p}weight`, [out, inputs]), read(`${p}bias`, [out])], {}, name);
+  const tokensOf = (map: string, channels: number, count: number): string =>
+    node('permute', [node('reshape', [map], { shape: [channels, count] })], { order: [1, 0] });
+  const channelNorm = (map: string, p: string, channels: number, rows: number): string =>
+    node(
+      'reshape',
+      [
+        node(
+          'permute',
+          [
+            node(
+              'layerNorm',
+              [
+                tokensOf(map, channels, rows * rows),
+                read(`${p}weight`, [channels]),
+                read(`${p}bias`, [channels]),
+              ],
+              { epsilon: CHANNEL_EPSILON },
+            ),
+          ],
+          { order: [1, 0] },
+        ),
+      ],
+      { shape: [channels, rows, rows] },
+    );
+  const perceptron = (x: string, p: string, sizes: readonly number[], name?: string): string => {
+    let y = x;
+    for (let i = 0; i + 1 < sizes.length; i += 1) {
+      const last = i + 2 === sizes.length;
+      y = linear(
+        y,
+        `${p}layers.${i}.`,
+        sizes[i + 1] as number,
+        sizes[i] as number,
+        last ? name : undefined,
+      );
+      if (i + 2 < sizes.length) y = node('relu', [y]);
+    }
+    return y;
+  };
+
+  /* The prompt encoder's weights, all of them: see the header. */
+  read(`${pe}pe_layer.positional_encoding_gaussian_matrix`, [2, dim / 2]);
+  for (let i = 0; i < 4; i += 1) read(`${pe}point_embeddings.${i}.weight`, [1, dim]);
+  read(`${pe}not_a_point_embed.weight`, [1, dim]);
+  const c = config.maskChannels;
+  const downscaled = (input: string): string => {
+    const p = `${pe}mask_downscaling.`;
+    let x = node(
+      'conv2d',
+      [input, read(`${p}0.weight`, [c / 4, 1, 2, 2]), read(`${p}0.bias`, [c / 4])],
+      { stride: 2 },
+    );
+    x = node('gelu', [channelNorm(x, `${p}1.`, c / 4, 2 * grid)]);
+    x = node('conv2d', [x, read(`${p}3.weight`, [c, c / 4, 2, 2]), read(`${p}3.bias`, [c])], {
+      stride: 2,
+    });
+    x = node('gelu', [channelNorm(x, `${p}4.`, c, grid)]);
+    x = node('conv2d', [x, read(`${p}6.weight`, [dim, c, 1, 1]), read(`${p}6.bias`, [dim])]);
+    return tokensOf(x, dim, cells);
+  };
+  const noMask = read(`${pe}no_mask_embed.weight`, [1, dim]);
+  if (mask === null) downscaledWeights(read, `${pe}mask_downscaling.`, dim, c);
+  const dense = mask === null ? node('reshape', [noMask], { shape: [dim] }) : downscaled(mask);
+
+  const positions = weights.derive(`${pe}grid_positions.${grid}x${grid}`, [cells, dim], (values) =>
+    samGridPositions(
+      values(`${pe}pe_layer.positional_encoding_gaussian_matrix`, [2, dim / 2]),
+      dim,
+      grid,
+    ),
+  );
+  const tokens = node(
+    'concat',
+    [
+      read(`${d}iou_token.weight`, [1, dim]),
+      read(`${d}mask_tokens.weight`, [masks, dim]),
+      ...(promptTokens === null ? [] : [promptTokens]),
+    ],
+    { axis: 0 },
+  );
+  const { queries, keys } = twoWay(
+    weights,
+    graph,
+    t,
+    config,
+    SAM_NAMES,
+    tokens,
+    node('add', [tokensOf(embedding, dim, cells), dense]),
+    positions,
+  );
+
+  /* Upsampled four times, and each mask token's vector dotted with every pixel. */
+  const u = `${d}output_upscaling.`;
+  const image = node('reshape', [node('permute', [keys], { order: [1, 0] })], {
+    shape: [dim, grid, grid],
+  });
+  let up = node(
+    'convTranspose2d',
+    [image, read(`${u}0.weight`, [dim, dim / 4, 2, 2]), read(`${u}0.bias`, [dim / 4])],
+    { stride: 2 },
+  );
+  up = node('gelu', [channelNorm(up, `${u}1.`, dim / 4, 2 * grid)]);
+  up = node('gelu', [
+    node(
+      'convTranspose2d',
+      [up, read(`${u}3.weight`, [dim / 4, dim / 8, 2, 2]), read(`${u}3.bias`, [dim / 8])],
+      { stride: 2 },
+    ),
+  ]);
+  const vectors = Array.from({ length: masks }, (_, i) =>
+    perceptron(
+      node('slice', [queries], { axis: 0, start: 1 + i, end: 2 + i }),
+      `${d}output_hypernetworks_mlps.${i}.`,
+      [dim, dim, dim, dim / 8],
+    ),
+  );
+  const pixels = 16 * cells;
+  const products = node('linear', [
+    tokensOf(up, dim / 8, pixels),
+    node('concat', vectors, { axis: 0 }),
+  ]);
+  node(
+    'reshape',
+    [node('permute', [products], { order: [1, 0] })],
+    { shape: [masks, 4 * grid, 4 * grid] },
+    names.masks,
+  );
+  perceptron(
+    node('slice', [queries], { axis: 0, start: 0, end: 1 }),
+    `${d}iou_prediction_head.`,
+    [dim, ...Array.from({ length: config.iouDepth - 1 }, () => config.iouHidden), masks],
+    names.quality,
+  );
+}
+
+/* The mask input's downscaling weights, read into a graph that takes no mask input. */
+function downscaledWeights(
+  read: (name: string, shape?: readonly number[]) => string,
+  p: string,
+  dim: number,
+  c: number,
+): void {
+  read(`${p}0.weight`, [c / 4, 1, 2, 2]);
+  read(`${p}0.bias`, [c / 4]);
+  read(`${p}1.weight`, [c / 4]);
+  read(`${p}1.bias`, [c / 4]);
+  read(`${p}3.weight`, [c, c / 4, 2, 2]);
+  read(`${p}3.bias`, [c]);
+  read(`${p}4.weight`, [c]);
+  read(`${p}4.bias`, [c]);
+  read(`${p}6.weight`, [dim, c, 1, 1]);
+  read(`${p}6.bias`, [dim]);
+}

package/src/models/samMiniature.ts

@@ -0,0 +1,190 @@
+/**
+ * A seeded miniature of MobileSAM: the upstream's layout, every tensor name and shape of it, at a
+ * size the CPU evaluates in a moment — what the model's tests and its hand-run oracle share, as the
+ * depth models' miniature is.
+ *
+ * **Small everywhere the upstream lets it be.** A 64-pixel image, four stages of 8, 16, 24 and 320
+ * channels — the last must be 320, since the upstream keeps its last merge's stride at one by that
+ * width alone — windows of 3, 4 and 2 over grids of 8, 4 and 4, so one stage pads its windows, one
+ * is attended whole and one divides evenly, and one or two blocks a stage. The neck's 256 channels
+ * are the upstream's own constant, so the decoder is 256 wide too; its MLP and quality head are
+ * narrow, and its grid is 4.
+ *
+ * **Every tensor the upstream holds is here**, the classification head the definition sets aside
+ * and the batch norms' counters included, so the oracle loads it strictly.
+ */
+import type { GraphTensor } from '@driftengine/texture';
+
+import type { MobileSamConfig } from './mobileSam.ts';
+import type { SamPrompt } from './samPrompt.ts';
+import { draw } from './seeded.ts';
+
+export const MINIATURE_MOBILE_SAM: MobileSamConfig = {
+  size: 64,
+  encoder: {
+    embedDims: [8, 16, 24, 320],
+    depths: [1, 2, 1, 1],
+    heads: [1, 2, 3, 10],
+    windows: [7, 3, 4, 2],
+    mergeStrides: [2, 2, 1],
+    mlpRatio: 4,
+    expandRatio: 4,
+    neck: 256,
+  },
+  decoder: {
+    dim: 256,
+    heads: 8,
+    mlp: 32,
+    depth: 2,
+    downsample: 2,
+    masks: 4,
+    iouHidden: 16,
+    iouDepth: 3,
+    maskChannels: 16,
+  },
+};
+
+/** The classes of the miniature's classification head, which nothing runs. */
+export const MINIATURE_SAM_CLASSES = 10;
+
+/**
+ * The original image the prompts are given in: 36 by 48, prepared at 48 by 64 in the encoder's 64
+ * square, so the masks are cropped before they come down to it and coordinates are scaled by 4/3.
+ */
+export const MINIATURE_SAM_IMAGE = { height: 36, width: 48 } as const;
+
+/** A point, which a padding point follows; a box; and two points with a box, which none follows. */
+export const MINIATURE_SAM_PROMPTS = {
+  point: { points: [[20.5, 14.25, 1]] },
+  box: { box: [6, 5, 40.5, 30] },
+  both: {
+    points: [
+      [20.5, 14.25, 1],
+      [33, 25.75, 0],
+    ],
+    box: [6, 5, 40.5, 30],
+  },
+} as const satisfies Record<string, SamPrompt>;
+
+/** Every tensor the upstream's encoder, prompt encoder and mask decoder hold. */
+export function miniatureSamCheckpoint(
+  config: MobileSamConfig,
+  seed: number,
+): Map<string, GraphTensor> {
+  const out = new Map<string, GraphTensor>();
+  const add = (name: string, shape: readonly number[]): void => {
+    out.set(name, { shape, data: draw(name, shape, seed) });
+  };
+  const convBn = (p: string, cout: number, cin: number, kernel: number): void => {
+    add(`${p}c.weight`, [cout, cin, kernel, kernel]);
+    for (const part of ['weight', 'bias', 'running_mean', 'running_var'])
+      add(`${p}bn.${part}`, [cout]);
+    out.set(`${p}bn.num_batches_tracked`, { shape: [], data: new Float32Array(1) });
+  };
+  const linear = (p: string, outputs: number, inputs: number): void => {
+    add(`${p}weight`, [outputs, inputs]);
+    add(`${p}bias`, [outputs]);
+  };
+  const norm = (p: string, dim: number): void => {
+    add(`${p}weight`, [dim]);
+    add(`${p}bias`, [dim]);
+  };
+
+  const { embedDims: dims, depths, windows, mlpRatio, expandRatio, neck } = config.encoder;
+  const e = 'image_encoder.';
+  convBn(`${e}patch_embed.seq.0.`, dims[0] / 2, 3, 3);
+  convBn(`${e}patch_embed.seq.2.`, dims[0], dims[0] / 2, 3);
+  for (let i = 0; i < depths[0]; i += 1) {
+    const p = `${e}layers.0.blocks.${i}.`;
+    const hidden = dims[0] * expandRatio;
+    convBn(`${p}conv1.`, hidden, dims[0], 1);
+    convBn(`${p}conv2.`, hidden, 1, 3);
+    convBn(`${p}conv3.`, dims[0], hidden, 1);
+  }
+  for (let stage = 0; stage < 4; stage += 1) {
+    const dim = dims[stage] as number;
+    if (stage > 0) {
+      for (let i = 0; i < (depths[stage] as number); i += 1) {
+        const p = `${e}layers.${stage}.blocks.${i}.`;
+        const window = windows[stage] as number;
+        add(`${p}attn.attention_biases`, [config.encoder.heads[stage] as number, window * window]);
+        norm(`${p}attn.norm.`, dim);
+        linear(`${p}attn.qkv.`, 3 * dim, dim);
+        linear(`${p}attn.proj.`, dim, dim);
+        norm(`${p}mlp.norm.`, dim);
+        linear(`${p}mlp.fc1.`, dim * mlpRatio, dim);
+        linear(`${p}mlp.fc2.`, dim, dim * mlpRatio);
+        convBn(`${p}local_conv.`, dim, 1, 3);
+      }
+    }
+    if (stage < 3) {
+      const next = dims[stage + 1] as number;
+      const p = `${e}layers.${stage}.downsample.`;
+      convBn(`${p}conv1.`, next, dim, 1);
+      convBn(`${p}conv2.`, next, 1, 3);
+      convBn(`${p}conv3.`, next, next, 1);
+    }
+  }
+  norm(`${e}norm_head.`, dims[3]);
+  linear(`${e}head.`, MINIATURE_SAM_CLASSES, dims[3]);
+  add(`${e}neck.0.weight`, [neck, dims[3], 1, 1]);
+  norm(`${e}neck.1.`, neck);
+  add(`${e}neck.2.weight`, [neck, neck, 3, 3]);
+  norm(`${e}neck.3.`, neck);
+
+  const { dim, masks, mlp, maskChannels: c, iouHidden, iouDepth } = config.decoder;
+  const pe = 'prompt_encoder.';
+  add(`${pe}pe_layer.positional_encoding_gaussian_matrix`, [2, dim / 2]);
+  for (let i = 0; i < 4; i += 1) add(`${pe}point_embeddings.${i}.weight`, [1, dim]);
+  add(`${pe}not_a_point_embed.weight`, [1, dim]);
+  add(`${pe}mask_downscaling.0.weight`, [c / 4, 1, 2, 2]);
+  add(`${pe}mask_downscaling.0.bias`, [c / 4]);
+  norm(`${pe}mask_downscaling.1.`, c / 4);
+  add(`${pe}mask_downscaling.3.weight`, [c, c / 4, 2, 2]);
+  add(`${pe}mask_downscaling.3.bias`, [c]);
+  norm(`${pe}mask_downscaling.4.`, c);
+  add(`${pe}mask_downscaling.6.weight`, [dim, c, 1, 1]);
+  add(`${pe}mask_downscaling.6.bias`, [dim]);
+  add(`${pe}no_mask_embed.weight`, [1, dim]);
+
+  const d = 'mask_decoder.';
+  const t = `${d}transformer.`;
+  const attention = (p: string, inner: number): void => {
+    for (const projection of ['q_proj', 'k_proj', 'v_proj'])
+      linear(`${p}${projection}.`, inner, dim);
+    linear(`${p}out_proj.`, dim, inner);
+  };
+  const inner = dim / config.decoder.downsample;
+  for (let i = 0; i < config.decoder.depth; i += 1) {
+    const p = `${t}layers.${i}.`;
+    attention(`${p}self_attn.`, dim);
+    norm(`${p}norm1.`, dim);
+    attention(`${p}cross_attn_token_to_image.`, inner);
+    norm(`${p}norm2.`, dim);
+    linear(`${p}mlp.lin1.`, mlp, dim);
+    linear(`${p}mlp.lin2.`, dim, mlp);
+    norm(`${p}norm3.`, dim);
+    norm(`${p}norm4.`, dim);
+    attention(`${p}cross_attn_image_to_token.`, inner);
+  }
+  attention(`${t}final_attn_token_to_image.`, inner);
+  norm(`${t}norm_final_attn.`, dim);
+  add(`${d}iou_token.weight`, [1, dim]);
+  add(`${d}mask_tokens.weight`, [masks, dim]);
+  add(`${d}output_upscaling.0.weight`, [dim, dim / 4, 2, 2]);
+  add(`${d}output_upscaling.0.bias`, [dim / 4]);
+  norm(`${d}output_upscaling.1.`, dim / 4);
+  add(`${d}output_upscaling.3.weight`, [dim / 4, dim / 8, 2, 2]);
+  add(`${d}output_upscaling.3.bias`, [dim / 8]);
+  for (let i = 0; i < masks; i += 1) {
+    const p = `${d}output_hypernetworks_mlps.${i}.layers.`;
+    linear(`${p}0.`, dim, dim);
+    linear(`${p}1.`, dim, dim);
+    linear(`${p}2.`, dim / 8, dim);
+  }
+  const sizes = [dim, ...Array.from({ length: iouDepth - 1 }, () => iouHidden), masks];
+  for (let i = 0; i + 1 < sizes.length; i += 1) {
+    linear(`${d}iou_prediction_head.layers.${i}.`, sizes[i + 1] as number, sizes[i] as number);
+  }
+  return out;
+}

package/src/models/samPrompt.ts

@@ -0,0 +1,227 @@
+/**
+ * What MobileSAM computes on the host rather than the device: a prompt's tokens, the positions of
+ * the embedding's grid, the size an image is prepared at, and a mask's stability.
+ *
+ * **A position is random Fourier features**: a point in the unit square, taken to [−1, 1], times a
+ * fixed Gaussian matrix of the checkpoint's, times 2π, and the sine then the cosine of each of those
+ * angles. The angles reach tens of radians, which the device's `sin` is only specified to 2^−11
+ * over [−π, π] for, so they are taken here, to single precision, with the engine's reproducible
+ * functions and in the upstream's order of roundings. That is also why a prompt is tokens given to
+ * the decoder rather than coordinates: a few tokens a prompt, against a grid of positions laid out
+ * once per size.
+ *
+ * **A prompt is the upstream's**: each point shifted to its pixel's centre and labelled on the
+ * object or off it, a padding point appended when there is no box, and a box as its two corners.
+ * Points and boxes are given in the original image's pixels and scaled to the prepared image as the
+ * upstream's predictor scales them.
+ */
+import { exactCos, exactSin } from '@driftengine/core';
+import type { WeightSource } from '@driftengine/texture';
+
+export interface SamPrompt {
+  /** Points in the original image's pixels, each labelled 1 on the object and 0 off it. */
+  readonly points?: readonly (readonly [x: number, y: number, label: 0 | 1])[];
+  /** A box in the original image's pixels: left, top, right, bottom. */
+  readonly box?: readonly [number, number, number, number];
+}
+
+const f = Math.fround;
+const TWO_PI = f(2 * Math.PI);
+
+/**
+ * Where a checkpoint keeps its prompt encoder, and how its processor places a prompt: SAM's and
+ * MobileSAM's, or Transformers' SAM 2.
+ */
+export interface SamPromptLayout {
+  /** The Fourier matrix, `[2, dim / 2]`. */
+  readonly gaussian: string;
+  /** A learned point embedding by kind — 0 off, 1 on, 2 and 3 a box's corners — as name and row. */
+  readonly embedding: (kind: number) => readonly [name: string, row: number];
+  readonly notAPoint: string;
+  /** Whether a box is followed by a padding point, as SAM 2's is. */
+  readonly boxPadding: boolean;
+  /**
+   * Whether a box's corners come first as points and a padding point always ends the prompt, as a
+   * SAM 2 video session orders them — where a prompt of nothing is a point meaning none, padded.
+   */
+  readonly boxAsPoints: boolean;
+  /**
+   * The prepared image's size for an original one: the longest side to `size` for SAM, the whole
+   * frame to a `size` square for SAM 2 — and whether a coordinate is scaled in double precision, as
+   * SAM's predictor does with numpy, or single, as SAM 2's processor does with a tensor.
+   */
+  readonly prepared: (height: number, width: number, size: number) => readonly [number, number];
+  readonly singleScale: boolean;
+}
+
+export const SAM_PROMPT: SamPromptLayout = {
+  gaussian: 'prompt_encoder.pe_layer.positional_encoding_gaussian_matrix',
+  embedding: (kind) => [`prompt_encoder.point_embeddings.${kind}.weight`, 0],
+  notAPoint: 'prompt_encoder.not_a_point_embed.weight',
+  boxPadding: false,
+  boxAsPoints: false,
+  prepared: (height, width, size) => samInputSize(height, width, size),
+  singleScale: false,
+};
+
+export const SAM2_PROMPT: SamPromptLayout = {
+  gaussian: 'prompt_encoder.shared_embedding.positional_embedding',
+  embedding: (kind) => ['prompt_encoder.point_embed.weight', kind],
+  notAPoint: 'prompt_encoder.not_a_point_embed.weight',
+  boxPadding: true,
+  boxAsPoints: false,
+  prepared: (_height, _width, size) => [size, size],
+  singleScale: true,
+};
+
+/** SAM 2 in a video session: the same encoder, a box as its two corner points first. */
+export const SAM2_VIDEO_PROMPT: SamPromptLayout = {
+  ...SAM2_PROMPT,
+  boxPadding: false,
+  boxAsPoints: true,
+};
+
+/** The upstream's `get_preprocess_shape`: the longest side to `longest`, each side rounded. */
+export function samInputSize(
+  height: number,
+  width: number,
+  longest: number,
+): readonly [number, number] {
+  const scale = (longest * 1.0) / Math.max(height, width);
+  return [Math.floor(height * scale + 0.5), Math.floor(width * scale + 0.5)];
+}
+
+/** How many tokens a prompt is, in a layout. */
+export function samTokenCount(request: SamPrompt, layout: SamPromptLayout = SAM_PROMPT): number {
+  const points = request.points?.length ?? 0;
+  if (layout.boxAsPoints) {
+    const given = points + (request.box === undefined ? 0 : 2);
+    return (given === 0 ? 1 : given) + 1;
+  }
+  const box = request.box === undefined ? 0 : layout.boxPadding ? 3 : 2;
+  return points > 0 ? points + (box === 0 ? 1 : box) : box;
+}
+
+/* Sines then cosines of a point in the unit square's angles, `2 · half` values from `at`. */
+function fourier(
+  gaussian: Float32Array,
+  half: number,
+  x: number,
+  y: number,
+  out: Float32Array,
+  at: number,
+): void {
+  const u = f(2 * x - 1);
+  const v = f(2 * y - 1);
+  for (let k = 0; k < half; k += 1) {
+    const angle = f(TWO_PI * f(u * (gaussian[k] as number) + v * (gaussian[half + k] as number)));
+    out[at + k] = exactSin(angle);
+    out[at + half + k] = exactCos(angle);
+  }
+}
+
+function weight(weights: WeightSource, name: string, length: number, row = 0): Float32Array {
+  const tensor = weights.get(name);
+  if (tensor === undefined || tensor.data.length < (row + 1) * length) {
+    throw new Error(`the prompt encoder has no "${name}" of ${length} values in row ${row}`);
+  }
+  return tensor.data.subarray(row * length, (row + 1) * length);
+}
+
+/**
+ * `[grid · grid, dim]`: the position of each cell of the embedding's grid, at its centre, as the
+ * upstream's `get_dense_pe` lays it out and the decoder reads it as tokens.
+ */
+export function samGridPositions(gaussian: Float32Array, dim: number, grid: number): Float32Array {
+  const out = new Float32Array(grid * grid * dim);
+  for (let y = 0; y < grid; y += 1) {
+    for (let x = 0; x < grid; x += 1) {
+      fourier(
+        gaussian,
+        dim / 2,
+        f((x + 0.5) / grid),
+        f((y + 0.5) / grid),
+        out,
+        (y * grid + x) * dim,
+      );
+    }
+  }
+  return out;
+}
+
+/**
+ * The tokens of `request`, `[samTokenCount(request, layout), dim]`, into `out`, for an original
+ * image of `height` by `width` prepared at `size` as `layout`'s processor prepares it.
+ */
+export function samPromptTokens(
+  weights: WeightSource,
+  dim: number,
+  size: number,
+  request: SamPrompt,
+  height: number,
+  width: number,
+  out: Float32Array,
+  layout: SamPromptLayout = SAM_PROMPT,
+): void {
+  const half = dim / 2;
+  const gaussian = weight(weights, layout.gaussian, 2 * half);
+  const [inputHeight, inputWidth] = layout.prepared(height, width, size);
+  const scaleX = layout.singleScale ? f(inputWidth / width) : inputWidth / width;
+  const scaleY = layout.singleScale ? f(inputHeight / height) : inputHeight / height;
+  const place = (value: number, scale: number): number =>
+    f(f(f(layout.singleScale ? f(value) * scale : value * scale) + 0.5) / size);
+  let at = 0;
+  /* A token: the position of a point at its pixel's centre, or none, and a learned embedding. */
+  const token = (x: number, y: number, positioned: boolean, name: string, row: number): void => {
+    if (positioned) {
+      fourier(gaussian, half, place(x, scaleX), place(y, scaleY), out, at);
+    } else {
+      out.fill(0, at, at + dim);
+    }
+    const learned = weight(weights, name, dim, row);
+    for (let c = 0; c < dim; c += 1) {
+      out[at + c] = f((out[at + c] as number) + (learned[c] as number));
+    }
+    at += dim;
+  };
+  if (layout.boxAsPoints) {
+    if (request.box !== undefined) {
+      const [left, top, right, bottom] = request.box;
+      token(left, top, true, ...layout.embedding(2));
+      token(right, bottom, true, ...layout.embedding(3));
+    }
+    for (const [x, y, label] of request.points ?? []) token(x, y, true, ...layout.embedding(label));
+    if (request.box === undefined && (request.points?.length ?? 0) === 0) {
+      token(0, 0, false, layout.notAPoint, 0);
+    }
+    token(0, 0, false, layout.notAPoint, 0);
+    return;
+  }
+  for (const [x, y, label] of request.points ?? []) token(x, y, true, ...layout.embedding(label));
+  if ((request.points?.length ?? 0) > 0 && request.box === undefined) {
+    token(0, 0, false, layout.notAPoint, 0);
+  }
+  if (request.box !== undefined) {
+    const [left, top, right, bottom] = request.box;
+    token(left, top, true, ...layout.embedding(2));
+    token(right, bottom, true, ...layout.embedding(3));
+    if (layout.boxPadding) token(0, 0, false, layout.notAPoint, 0);
+  }
+}
+
+/**
+ * The upstream's stability score: the share of the pixels above `threshold − offset` that are also
+ * above `threshold + offset` — one over the other of the mask thresholded high and thresholded low,
+ * since one always lies inside the other. A mask with no pixel above the lower threshold scores
+ * NaN, as the upstream's does.
+ */
+export function stabilityScore(logits: Float32Array, threshold: number, offset: number): number {
+  let inside = 0;
+  let union = 0;
+  for (let i = 0; i < logits.length; i += 1) {
+    const value = logits[i] as number;
+    if (value > threshold + offset) inside += 1;
+    if (value > threshold - offset) union += 1;
+  }
+  return inside / union;
+}

package/src/models/seeded.ts

@@ -0,0 +1,50 @@
+/**
+ * A miniature checkpoint's values, drawn from a seed and a tensor's name, so each tensor depends on
+ * its name and never on the order the tensors were made in.
+ */
+import { mulberry32 } from '@driftengine/core';
+
+/* FNV-1a over a name, so each tensor's values depend on its name and not on the order made. */
+export function nameSeed(name: string, seed: number): number {
+  let hash = 0x811c9dc5 ^ seed;
+  for (let i = 0; i < name.length; i += 1) {
+    hash = Math.imul(hash ^ name.charCodeAt(i), 0x01000193);
+  }
+  return hash >>> 0;
+}
+
+/**
+ * The spread a tensor's values are drawn with, by what the name says it is: a norm's scale near
+ * one, a layer scale near a half, a bias small, and a weight at 1.7/√fan-in. At 1/√fan-in the
+ * head's twenty convolutions shrank the signal until the depth was its biases — 1.08 to 1.11
+ * everywhere, and a slip in a norm's ε moved nothing a test could see; at He's √(6/fan-in) the
+ * fusion's residual sums compounded it to depths of 10^33. At 1.7 the depth spans 0.96 to 11.7.
+ *
+ * **A one-dimensional weight is a norm's scale** in every checkpoint here, whatever its name — a
+ * channel norm is often a bare index in a sequence — and is drawn near one.
+ *
+ * **A batch norm's statistics are drawn as statistics**: a running variance between a half and one
+ * and a half, so the scale folded from it stays near one, and a running mean small. Drawn as any
+ * other vector, a variance near zero would fold into a scale of hundreds.
+ */
+export function draw(name: string, shape: readonly number[], seed: number): Float32Array {
+  const next = mulberry32(nameSeed(name, seed));
+  const count = shape.reduce((total, d) => total * d, 1);
+  const fanIn = shape.slice(1).reduce((total, d) => total * d, 1);
+  const norm = (/norm|\.bn\./.test(name) || shape.length === 1) && name.endsWith('.weight');
+  const gamma = name.endsWith('.gamma') || name.endsWith('.lambda1');
+  const variance = name.endsWith('.running_var');
+  const centre = norm || variance ? 1 : gamma ? 0.5 : 0;
+  const spread = norm
+    ? 0.1
+    : gamma
+      ? 0.2
+      : variance
+        ? 0.5
+        : name.endsWith('.bias') || name.endsWith('.running_mean')
+          ? 0.1
+          : shape.length >= 2 && name.endsWith('.weight')
+            ? 1.7 / Math.sqrt(fanIn)
+            : 0.5;
+  return Float32Array.from({ length: count }, () => centre + (next() * 2 - 1) * spread);
+}