@driftengine/capture 4.0.0

package/src/models/sam21Video.ts

@@ -0,0 +1,154 @@
+/**
+ * The graphs a SAM 2.1 video adds to its image graphs: the memory encoder, the memory attention for
+ * each count of remembered frames and pointer tokens, and the mask brought to the frame's size for
+ * the memory. `Sam21Tracker` runs them.
+ */
+import type { Architecture } from '@driftengine/texture';
+
+import type { Sam21Config } from './sam21.ts';
+import { sam2MemoryAttention, sam2MemoryEncoder } from './sam2Memory.ts';
+
+/* Everything but the memory encoder's weights, for a graph that reads only them. */
+const IMAGE = [
+  'vision_encoder.',
+  'prompt_encoder.',
+  'mask_decoder.',
+  'shared_image_embedding.',
+  'object_pointer_proj.',
+  'memory_attention.',
+  'mask_downsample.',
+  'temporal_positional_encoding_projection_layer.',
+  'memory_temporal_positional_encoding',
+  'no_memory_embedding',
+  'no_memory_positional_encoding',
+  'no_object_pointer',
+] as const;
+
+/**
+ * `features`, `[dim, grid, grid]`, and `mask`, `[1, size, size]` — logits, or a binary mask when
+ * `binary`, as a frame prompted with points is encoded — to `memory`, `[memory.dim, grid, grid]`.
+ * The host adds the occlusion embedding where the object is absent, and keeps the memory rounded to
+ * bfloat16, as the upstream keeps it; this graph reads the embedding so the file holds it.
+ */
+export function sam21MemoryEncoder(config: Sam21Config, binary: boolean): Architecture {
+  return (weights, graph) => {
+    for (const name of IMAGE) weights.ignore(name);
+    const { dim } = config.decoder;
+    const grid = config.size / 16;
+    weights.read('occlusion_spatial_embedding_parameter', [1, config.memory.dim]);
+    sam2MemoryEncoder(
+      weights,
+      graph,
+      dim,
+      config.memory,
+      grid,
+      'features',
+      'mask',
+      binary,
+      'memory',
+    );
+    return {
+      inputs: [
+        { name: 'features', shape: [dim, grid, grid] },
+        { name: 'mask', shape: [1, config.size, config.size] },
+      ],
+      outputs: ['memory'],
+    };
+  };
+}
+
+/**
+ * `features`, `[dim, grid, grid]`, `memory` and its `positions`, `[frames·grid² + pointers,
+ * memory.dim]`, to `conditioned`, `[dim, grid, grid]`: one graph for each count of frames and of
+ * pointer tokens a tracker holds. The host makes the memory's positions, so this graph reads the
+ * temporal table and the pointers' projection for the file to hold them.
+ */
+export function sam21MemoryAttention(
+  config: Sam21Config,
+  frames: number,
+  pointers: number,
+): Architecture {
+  return (weights, graph) => {
+    for (const name of [
+      'vision_encoder.',
+      'prompt_encoder.',
+      'mask_decoder.',
+      'shared_image_embedding.',
+      'object_pointer_proj.',
+      'memory_encoder.',
+      'mask_downsample.',
+      'no_memory_embedding',
+      'no_memory_positional_encoding',
+      'no_object_pointer',
+      'occlusion_spatial_embedding_parameter',
+    ]) {
+      weights.ignore(name);
+    }
+    const { dim } = config.decoder;
+    const grid = config.size / 16;
+    const m = config.memory;
+    weights.read('memory_temporal_positional_encoding', [m.frames, 1, 1, m.dim]);
+    weights.read('temporal_positional_encoding_projection_layer.weight', [m.dim, dim]);
+    weights.read('temporal_positional_encoding_projection_layer.bias', [m.dim]);
+    const tokens = frames * grid * grid + pointers;
+    sam2MemoryAttention(
+      weights,
+      graph,
+      dim,
+      m,
+      grid,
+      frames,
+      pointers,
+      'features',
+      'memory',
+      'positions',
+      'conditioned',
+    );
+    return {
+      inputs: [
+        { name: 'features', shape: [dim, grid, grid] },
+        { name: 'memory', shape: [tokens, m.dim] },
+        { name: 'positions', shape: [tokens, m.dim] },
+      ],
+      outputs: ['conditioned'],
+    };
+  };
+}
+
+/** `masks`, `[count, 4·grid, 4·grid]`, to `high`, `[count, size, size]`, bilinearly. Reads no weights. */
+export function sam21Upscale(config: Sam21Config, count: number): Architecture {
+  return (_weights, graph) => {
+    const low = config.size / 4;
+    graph.node(
+      'resize',
+      ['masks'],
+      { height: config.size, width: config.size, mode: 'bilinear', alignCorners: false },
+      'high',
+    );
+    return { inputs: [{ name: 'masks', shape: [count, low, low] }], outputs: ['high'] };
+  };
+}
+
+/**
+ * `masks`, `[count, size / 4, size / 4]`, to `image`, `[count, height, width]`: the tracker's masks
+ * brought to the original frame bilinearly, as the upstream's `post_process_masks` brings them —
+ * straight from a quarter of the square, since SAM 2 resizes a frame without padding it. Reads no
+ * weights.
+ */
+export function sam21MasksToImage(
+  config: Sam21Config,
+  count: number,
+  height: number,
+  width: number,
+): Architecture {
+  return (_weights, graph) => {
+    const low = config.size / 4;
+    graph.node(
+      'resize',
+      ['masks'],
+      { height, width, mode: 'bilinear', alignCorners: false },
+      'image',
+    );
+    return { inputs: [{ name: 'masks', shape: [count, low, low] }], outputs: ['image'] };
+  };
+}

package/src/models/sam2Choice.ts

@@ -0,0 +1,62 @@
+/**
+ * Which of a frame's four decoded masks a SAM 2 tracker keeps, and which pointer goes with it, as the
+ * upstream's video decoder chooses.
+ *
+ * **A prompt of one point or none keeps the best of the three ambiguous masks** by predicted quality,
+ * and that mask token's pointer. **More points keep the single mask**, unless its stability — pixels
+ * above +0.05 over pixels above −0.05, one where there are none — is under 0.98, when the best of the
+ * three is kept instead; **the pointer is the single mask's either way**, as the upstream slices its
+ * tokens before it falls back. **An object whose score is not positive is absent**: its mask is −1024
+ * everywhere and its pointer the learned one for no object. A box counts as its two corner points.
+ */
+
+const NO_OBJECT = -1024;
+const STABILITY_DELTA = 0.05;
+const STABILITY_THRESHOLD = 0.98;
+
+/**
+ * The kept mask's logits and pointer, from `masks`, `[4, cells]`, `quality`, `[4]`, `pointers`,
+ * `[4, dim]`, the object's score, the prompt's point count and the no-object pointer.
+ */
+export function chooseMask(
+  masks: Float32Array,
+  quality: Float32Array,
+  pointers: Float32Array,
+  score: number,
+  points: number,
+  noObject: Float32Array,
+): { readonly lowRes: Float32Array; readonly pointer: Float32Array } {
+  const count = quality.length;
+  const cells = masks.length / count;
+  const dim = pointers.length / count;
+  let best = 1;
+  for (let m = 2; m < count; m += 1)
+    if ((quality[m] as number) > (quality[best] as number)) best = m;
+  let mask = best;
+  let token = best;
+  if (points > 1) {
+    mask = stability(masks.subarray(0, cells)) >= STABILITY_THRESHOLD ? 0 : best;
+    token = 0;
+  }
+  if (score <= 0) {
+    return {
+      lowRes: new Float32Array(cells).fill(NO_OBJECT),
+      pointer: Float32Array.from(noObject),
+    };
+  }
+  return {
+    lowRes: Float32Array.from(masks.subarray(mask * cells, (mask + 1) * cells)),
+    pointer: Float32Array.from(pointers.subarray(token * dim, (token + 1) * dim)),
+  };
+}
+
+/* Pixels above +δ over pixels above −δ, and one where none are. */
+function stability(logits: Float32Array): number {
+  let inside = 0;
+  let union = 0;
+  for (let i = 0; i < logits.length; i += 1) {
+    if ((logits[i] as number) > STABILITY_DELTA) inside += 1;
+    if ((logits[i] as number) > -STABILITY_DELTA) union += 1;
+  }
+  return union > 0 ? inside / union : 1;
+}

package/src/models/sam2Decoder.ts

@@ -0,0 +1,218 @@
+/**
+ * SAM 2's mask decoder: SAM's two-way transformer with the encoder's finer levels brought back in,
+ * an object score, and every mask token's object pointer for a video's memory.
+ *
+ * **The tokens lead with an object-score token**, then the quality token and one per mask, then
+ * the prompt's. **The upsampling adds the encoder's skips**: its second level after the first
+ * transposed convolution, and its first after the second, each before the activation. Quality is
+ * a sigmoid, as the upstream's `sigmoid_output`; the object score is a logit, positive where the
+ * object is present. **Each mask token's object pointer is projected here**, all four, so the host
+ * can keep the one it chooses without another graph. Norms of channels take ε 1e-6.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+import type { SamDecoderConfig } from './samDecoder.ts';
+import { samGridPositions } from './samPrompt.ts';
+import { SAM2_NAMES, twoWay } from './twoWay.ts';
+
+const CHANNEL_EPSILON = 1e-6;
+
+export interface Sam2DecoderInputs {
+  /** `[dim, grid, grid]`: the features, conditioned on memory or given `no_memory_embedding`. */
+  readonly embedding: string;
+  /** `[tokens, dim]`, or null for a prompt of none. */
+  readonly prompt: string | null;
+  /** `[1, 4·grid, 4·grid]`: a mask's logits to refine, or null. */
+  readonly mask: string | null;
+  /** The encoder's narrowed finer levels, `[dim / 4, 2·grid, 2·grid]` and `[dim / 8, 4·grid, 4·grid]`. */
+  readonly high1: string;
+  readonly high0: string;
+}
+
+/** What the decoder's four outputs are named. */
+export interface Sam2DecoderOutput {
+  /** `[masks, 4·grid, 4·grid]` logits. */
+  readonly masks: string;
+  /** `[1, masks]`, each mask's predicted intersection over union. */
+  readonly quality: string;
+  /** `[1, 1]`: the object score's logit. */
+  readonly object: string;
+  /** `[masks, dim]`: each mask token's object pointer. */
+  readonly pointers: string;
+}
+
+export function sam2Decoder(
+  weights: Weights,
+  graph: GraphBuilder,
+  config: SamDecoderConfig,
+  grid: number,
+  inputs: Sam2DecoderInputs,
+  names: Sam2DecoderOutput,
+): void {
+  const node = graph.node.bind(graph);
+  const read = weights.read.bind(weights);
+  const { dim, masks } = config;
+  const cells = grid * grid;
+  const m = 'mask_decoder.';
+  const pe = 'prompt_encoder.';
+  const linear = (x: string, p: string, out: number, width: number, name?: string): string =>
+    node('linear', [x, read(`${p}weight`, [out, width]), read(`${p}bias`, [out])], {}, name);
+  /* The upstream's perceptron: `proj_in`, then `layers`, then `proj_out`, ReLU between. */
+  const perceptron = (x: string, p: string, sizes: readonly number[], name?: string): string => {
+    const parts = sizes.length - 1;
+    let y = x;
+    for (let i = 0; i < parts; i += 1) {
+      const at = i === 0 ? 'proj_in' : i === parts - 1 ? 'proj_out' : `layers.${i - 1}`;
+      const last = i === parts - 1;
+      y = linear(
+        y,
+        `${p}${at}.`,
+        sizes[i + 1] as number,
+        sizes[i] as number,
+        last ? name : undefined,
+      );
+      if (!last) y = node('relu', [y]);
+    }
+    return y;
+  };
+  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 => {
+    const normed = node(
+      'layerNorm',
+      [
+        tokensOf(map, channels, rows * rows),
+        read(`${p}weight`, [channels]),
+        read(`${p}bias`, [channels]),
+      ],
+      { epsilon: CHANNEL_EPSILON },
+    );
+    return node('reshape', [node('permute', [normed], { order: [1, 0] })], {
+      shape: [channels, rows, rows],
+    });
+  };
+
+  /* The prompt encoder's weights, all of them, for the host's tokens and the graph that refines. */
+  read(`${pe}shared_embedding.positional_embedding`, [2, dim / 2]);
+  read(`${pe}point_embed.weight`, [4, dim]);
+  read(`${pe}not_a_point_embed.weight`, [1, dim]);
+  const c = config.maskChannels;
+  const embed = `${pe}mask_embed.`;
+  const convs = [
+    [`${embed}conv1.`, c / 4, 1, 2],
+    [`${embed}conv2.`, c, c / 4, 2],
+    [`${embed}conv3.`, dim, c, 1],
+  ] as const;
+  const w = convs.map(([p, out, into, k]) => [
+    read(`${p}weight`, [out, into, k, k]),
+    read(`${p}bias`, [out]),
+  ]);
+  let dense = node('reshape', [read(`${pe}no_mask_embed.weight`, [1, dim])], { shape: [dim] });
+  if (inputs.mask !== null) {
+    let x = node('conv2d', [inputs.mask, ...(w[0] as string[])], { stride: 2 });
+    x = node('gelu', [channelNorm(x, `${embed}layer_norm1.`, c / 4, 2 * grid)]);
+    x = node('conv2d', [x, ...(w[1] as string[])], { stride: 2 });
+    x = node('gelu', [channelNorm(x, `${embed}layer_norm2.`, c, grid)]);
+    dense = tokensOf(node('conv2d', [x, ...(w[2] as string[])]), dim, cells);
+  } else {
+    for (const [i, channels] of [
+      [1, c / 4],
+      [2, c],
+    ] as const) {
+      read(`${embed}layer_norm${i}.weight`, [channels]);
+      read(`${embed}layer_norm${i}.bias`, [channels]);
+    }
+  }
+
+  const positions = weights.derive(
+    `shared_image_embedding.grid_positions.${grid}x${grid}`,
+    [cells, dim],
+    (values) =>
+      samGridPositions(
+        values('shared_image_embedding.positional_embedding', [2, dim / 2]),
+        dim,
+        grid,
+      ),
+  );
+  read('shared_image_embedding.positional_embedding', [2, dim / 2]);
+  const tokens = node(
+    'concat',
+    [
+      read(`${m}obj_score_token.weight`, [1, dim]),
+      read(`${m}iou_token.weight`, [1, dim]),
+      read(`${m}mask_tokens.weight`, [masks, dim]),
+      ...(inputs.prompt === null ? [] : [inputs.prompt]),
+    ],
+    { axis: 0 },
+  );
+  const image = node('add', [tokensOf(inputs.embedding, dim, cells), dense]);
+  const { queries, keys } = twoWay(
+    weights,
+    graph,
+    `${m}transformer.`,
+    config,
+    SAM2_NAMES,
+    tokens,
+    image,
+    positions,
+  );
+
+  /* Upsampled four times with the encoder's skips, and each mask token's vector dotted with every pixel. */
+  const map = node('reshape', [node('permute', [keys], { order: [1, 0] })], {
+    shape: [dim, grid, grid],
+  });
+  let up = node(
+    'convTranspose2d',
+    [
+      map,
+      read(`${m}upscale_conv1.weight`, [dim, dim / 4, 2, 2]),
+      read(`${m}upscale_conv1.bias`, [dim / 4]),
+    ],
+    { stride: 2 },
+  );
+  up = node('gelu', [
+    channelNorm(node('add', [up, inputs.high1]), `${m}upscale_layer_norm.`, dim / 4, 2 * grid),
+  ]);
+  up = node(
+    'convTranspose2d',
+    [
+      up,
+      read(`${m}upscale_conv2.weight`, [dim / 4, dim / 8, 2, 2]),
+      read(`${m}upscale_conv2.bias`, [dim / 8]),
+    ],
+    { stride: 2 },
+  );
+  up = node('gelu', [node('add', [up, inputs.high0])]);
+  const maskTokens = node('slice', [queries], { axis: 0, start: 2, end: 2 + masks });
+  const vectors = Array.from({ length: masks }, (_, i) =>
+    perceptron(
+      node('slice', [maskTokens], { axis: 0, start: i, end: i + 1 }),
+      `${m}output_hypernetworks_mlps.${i}.`,
+      [dim, dim, dim, dim / 8],
+    ),
+  );
+  const products = node('linear', [
+    tokensOf(up, dim / 8, 16 * cells),
+    node('concat', vectors, { axis: 0 }),
+  ]);
+  node(
+    'reshape',
+    [node('permute', [products], { order: [1, 0] })],
+    { shape: [masks, 4 * grid, 4 * grid] },
+    names.masks,
+  );
+  const hidden = Array.from({ length: config.iouDepth - 1 }, () => config.iouHidden);
+  const quality = perceptron(
+    node('slice', [queries], { axis: 0, start: 1, end: 2 }),
+    `${m}iou_prediction_head.`,
+    [dim, ...hidden, masks],
+  );
+  node('sigmoid', [quality], {}, names.quality);
+  perceptron(
+    node('slice', [queries], { axis: 0, start: 0, end: 1 }),
+    `${m}pred_obj_score_head.`,
+    [dim, dim, dim, 1],
+    names.object,
+  );
+  perceptron(maskTokens, 'object_pointer_proj.', [dim, dim, dim, dim], names.pointers);
+}

package/src/models/sam2Memories.ts

@@ -0,0 +1,165 @@
+/**
+ * What a SAM 2 tracker keeps between frames, and how it chooses what a new frame reads: the host's
+ * half of the video memory, as the upstream's `Sam2VideoModel` does it.
+ *
+ * **A frame reads every prompted frame's memory and the six frames before it** — after it, tracking
+ * backwards — each at its distance, which picks a row of a learned temporal table; a prompted frame
+ * sits at distance zero, which the upstream indexes as the table's last row. **And the object
+ * pointers** of the prompted frames up to it and of the frames up to fifteen before it, each placed
+ * by a sine of its distance over the most there could be, projected to the memory's width; a
+ * pointer wider than a memory token is cut into as many tokens as it takes.
+ *
+ * **A memory is kept at bfloat16**, rounded to nearest even as the upstream stores it; what the
+ * tracker holds is the rounded value, so a frame reads what the upstream's reads.
+ */
+import type { WeightSource } from '@driftengine/texture';
+
+import { sam2PointerPositions } from './sam2Positions.ts';
+
+/** One object's output on one frame, as the tracker keeps it. */
+export interface Sam2Output {
+  /** `[(size / 4)²]`: the mask's logits, −1024 everywhere when the object is absent. */
+  readonly lowRes: Float32Array;
+  /** `[dim]`: the object pointer, or the learned one for an absent object. */
+  readonly pointer: Float32Array;
+  readonly score: number;
+  /** `[cells, memory.dim]` rounded to bfloat16, once the memory is encoded. */
+  memory: Float32Array | null;
+}
+
+/** One object's history: its prompted frames' outputs and the frames it was tracked on. */
+export interface Sam2History {
+  readonly prompted: Map<number, Sam2Output>;
+  readonly tracked: Map<number, Sam2Output>;
+}
+
+/** Rounds every value to bfloat16, to nearest even, in place. */
+export function toBfloat16(values: Float32Array): void {
+  const bits = new Uint32Array(values.buffer, values.byteOffset, values.length);
+  for (let i = 0; i < bits.length; i += 1) {
+    const word = bits[i] as number;
+    if ((word & 0x7f800000) === 0x7f800000) {
+      bits[i] = (word & 0xffff0000) | ((word & 0xffff) === 0 ? 0 : 0x400000);
+      continue;
+    }
+    bits[i] = ((word + 0x7fff + ((word >>> 16) & 1)) & 0xffff0000) >>> 0;
+  }
+}
+
+/** The memories a frame reads, each with its distance, prompted frames first. */
+export function memoriesFor(
+  history: Sam2History,
+  frame: number,
+  frames: number,
+  reverse: boolean,
+): readonly (readonly [distance: number, output: Sam2Output])[] {
+  const chosen: (readonly [number, Sam2Output])[] = [...history.prompted.values()].map(
+    (output) => [0, output] as const,
+  );
+  for (let distance = frames - 1; distance > 0; distance -= 1) {
+    const output = history.tracked.get(reverse ? frame + distance : frame - distance);
+    if (output !== undefined) chosen.push([distance, output]);
+  }
+  return chosen.filter(([, output]) => output.memory !== null);
+}
+
+/** The object pointers a frame reads, each with its distance. */
+export function pointersFor(
+  history: Sam2History,
+  frame: number,
+  videoFrames: number,
+  most: number,
+  reverse: boolean,
+): {
+  readonly distances: readonly number[];
+  readonly pointers: readonly Float32Array[];
+  readonly span: number;
+} {
+  const span = Math.min(videoFrames, most);
+  const distances: number[] = [];
+  const pointers: Float32Array[] = [];
+  for (const [at, output] of history.prompted) {
+    if (reverse ? at >= frame : at <= frame) {
+      distances.push((frame - at) * (reverse ? -1 : 1));
+      pointers.push(output.pointer);
+    }
+  }
+  for (let distance = 1; distance < span; distance += 1) {
+    const at = reverse ? frame + distance : frame - distance;
+    if (at < 0 || at >= videoFrames) break;
+    const output = history.tracked.get(at);
+    if (output !== undefined) {
+      distances.push(distance);
+      pointers.push(output.pointer);
+    }
+  }
+  return { distances, pointers, span };
+}
+
+/**
+ * The memory tokens and their positions a frame's memory attention reads: the frames' cells, each
+ * frame's positions its cells' sines plus its distance's temporal row, then the pointers cut into
+ * tokens of `width`, each carrying its pointer's projected sine position.
+ */
+export function memoryTokens(
+  weights: WeightSource,
+  memories: readonly (readonly [number, Sam2Output])[],
+  pointers: {
+    readonly distances: readonly number[];
+    readonly pointers: readonly Float32Array[];
+    readonly span: number;
+  },
+  cellPositions: Float32Array,
+  width: number,
+  dim: number,
+): {
+  readonly memory: Float32Array;
+  readonly positions: Float32Array;
+  readonly pointerTokens: number;
+} {
+  const temporal = tensor(weights, 'memory_temporal_positional_encoding');
+  const cells = cellPositions.length / width;
+  const rows = temporal.length / width;
+  const splits = dim / width;
+  const pointerTokens = pointers.pointers.length * splits;
+  const total = memories.length * cells + pointerTokens;
+  const memory = new Float32Array(total * width);
+  const positions = new Float32Array(total * width);
+  memories.forEach(([distance, output], m) => {
+    memory.set(output.memory as Float32Array, m * cells * width);
+    const row = ((distance - 1 + rows) % rows) * width;
+    for (let i = 0; i < cells * width; i += 1) {
+      positions[m * cells * width + i] = Math.fround(
+        (cellPositions[i] as number) + (temporal[row + (i % width)] as number),
+      );
+    }
+  });
+  if (pointerTokens > 0) {
+    const at = memories.length * cells * width;
+    const sine = sam2PointerPositions(
+      pointers.distances.map((distance) => Math.fround(distance / Math.fround(pointers.span - 1))),
+      dim,
+    );
+    const weight = tensor(weights, 'temporal_positional_encoding_projection_layer.weight');
+    const bias = tensor(weights, 'temporal_positional_encoding_projection_layer.bias');
+    pointers.pointers.forEach((pointer, p) => {
+      for (let s = 0; s < splits; s += 1) {
+        const token = at + (p * splits + s) * width;
+        memory.set(pointer.subarray(s * width, (s + 1) * width), token);
+        for (let o = 0; o < width; o += 1) {
+          let sum = bias[o] as number;
+          for (let c = 0; c < dim; c += 1)
+            sum += (weight[o * dim + c] as number) * (sine[p * dim + c] as number);
+          positions[token + o] = Math.fround(sum);
+        }
+      }
+    });
+  }
+  return { memory, positions, pointerTokens };
+}
+
+function tensor(weights: WeightSource, name: string): Float32Array {
+  const held = weights.get(name);
+  if (held === undefined) throw new Error(`the memory has no "${name}"`);
+  return held.data;
+}