@driftengine/capture 4.0.0

package/src/models/sam2Memory.ts

@@ -0,0 +1,247 @@
+/**
+ * SAM 2's memory: the encoder that makes a frame and its mask into a memory, and the attention that
+ * conditions a new frame's features on the memories kept.
+ *
+ * **A memory is the frame's features where the mask says the object is**: the mask, taken through a
+ * sigmoid unless it is already binary, scaled to ±10 and downsampled sixteen times by strided 3×3
+ * convolutions, is added to the features projected once, and the sum is fused by ConvNeXt blocks —
+ * a depthwise 7×7, a channel norm, and a pointwise MLP scaled per channel — and narrowed to the
+ * memory's 64 channels. Norms of channels take ε 1e-6.
+ *
+ * **The attention reads every memory at once**: the current features, with a tenth of their sine
+ * positions added, attend to themselves and then to the memories' tokens — every frame's cells and
+ * the object pointers after them — then pass an MLP, four times. Queries and the frames' keys are
+ * turned by an axial rotary embedding of the cell's column and row, the same table for every frame;
+ * the pointers' keys are not turned, as the upstream's `num_k_exclude_rope` leaves them. Its norms
+ * take PyTorch's default ε, 1e-5.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+import type { Sam21MemoryConfig } from './sam21.ts';
+import { sam2Rotary, sam2SinePositions } from './sam2Positions.ts';
+
+const CHANNEL_EPSILON = 1e-6;
+const EPSILON = 1e-5;
+
+/**
+ * `[memory.dim, grid, grid]` into `output`, from `features`, `[dim, grid, grid]`, and `mask`,
+ * `[1, 16·grid, 16·grid]` — logits, or a binary mask when `binary`.
+ */
+export function sam2MemoryEncoder(
+  weights: Weights,
+  graph: GraphBuilder,
+  dim: number,
+  memory: Sam21MemoryConfig,
+  grid: number,
+  features: string,
+  mask: string,
+  binary: boolean,
+  output: string,
+): void {
+  const node = graph.node.bind(graph);
+  const read = weights.read.bind(weights);
+  const e = 'memory_encoder.';
+  const side = 16 * grid;
+  const conv = (
+    x: string,
+    p: string,
+    out: number,
+    into: number,
+    k: number,
+    attributes = {},
+  ): string =>
+    node('conv2d', [x, read(`${p}weight`, [out, into, k, k]), read(`${p}bias`, [out])], attributes);
+  const tokensOf = (map: string, channels: number, count: number): string =>
+    node('permute', [node('reshape', [map], { shape: [channels, count] })], { order: [1, 0] });
+  const mapOf = (tokens: string, channels: number, rows: number): string =>
+    node('reshape', [node('permute', [tokens], { order: [1, 0] })], {
+      shape: [channels, rows, rows],
+    });
+  const channelNorm = (map: string, p: string, channels: number, rows: number): string =>
+    mapOf(
+      node(
+        'layerNorm',
+        [
+          tokensOf(map, channels, rows * rows),
+          read(`${p}weight`, [channels]),
+          read(`${p}bias`, [channels]),
+        ],
+        { epsilon: CHANNEL_EPSILON },
+      ),
+      channels,
+      rows,
+    );
+
+  /* The mask to ±10: a sigmoid's 0 to 1, or a binary mask's, times 20 less 10. */
+  const unit = binary ? mask : node('sigmoid', [mask]);
+  const scale = weights.constant(
+    `sam2.mask_scale.${side}`,
+    [side],
+    new Float32Array(side).fill(20),
+  );
+  const shift = weights.constant(
+    `sam2.mask_shift.${side}`,
+    [side],
+    new Float32Array(side).fill(-10),
+  );
+  let down = node('add', [node('mul', [unit, scale]), shift]);
+  let channels = 1;
+  let rows = side;
+  for (let i = 0; i < 4; i += 1) {
+    const p = `${e}mask_downsampler.layers.${i}.`;
+    down = conv(down, `${p}conv.`, channels * 4, channels, 3, { stride: 2, padding: 1 });
+    channels *= 4;
+    rows /= 2;
+    down = node('gelu', [channelNorm(down, `${p}layer_norm.`, channels, rows)]);
+  }
+  down = conv(down, `${e}mask_downsampler.final_conv.`, dim, channels, 1);
+  let fused = node('add', [conv(features, `${e}feature_projection.`, dim, dim, 1), down]);
+  for (let i = 0; i < memory.fuserLayers; i += 1) {
+    const p = `${e}memory_fuser.layers.${i}.`;
+    const spread = conv(fused, `${p}depthwise_conv.`, dim, 1, 7, { padding: 3, groups: dim });
+    const normed = tokensOf(channelNorm(spread, `${p}layer_norm.`, dim, grid), dim, grid * grid);
+    const hidden = node('gelu', [
+      node('linear', [
+        normed,
+        read(`${p}pointwise_conv1.weight`, [memory.fuserHidden, dim]),
+        read(`${p}pointwise_conv1.bias`, [memory.fuserHidden]),
+      ]),
+    ]);
+    const back = node('linear', [
+      hidden,
+      read(`${p}pointwise_conv2.weight`, [dim, memory.fuserHidden]),
+      read(`${p}pointwise_conv2.bias`, [dim]),
+    ]);
+    fused = node('add', [fused, mapOf(node('mul', [back, read(`${p}scale`, [dim])]), dim, grid)]);
+  }
+  node(
+    'conv2d',
+    [
+      fused,
+      read(`${e}projection.weight`, [memory.dim, dim, 1, 1]),
+      read(`${e}projection.bias`, [memory.dim]),
+    ],
+    {},
+    output,
+  );
+}
+
+/**
+ * `[dim, grid, grid]` into `output`: `features`, `[dim, grid, grid]`, conditioned on `memory`,
+ * `[frames·grid² + pointers, memory.dim]`, whose positions are `positions` of that shape — the
+ * frames' cells first and `pointers` object-pointer tokens after them.
+ */
+export function sam2MemoryAttention(
+  weights: Weights,
+  graph: GraphBuilder,
+  dim: number,
+  config: Sam21MemoryConfig,
+  grid: number,
+  frames: number,
+  pointers: number,
+  features: string,
+  memory: string,
+  positions: string,
+  output: string,
+): void {
+  const node = graph.node.bind(graph);
+  const read = weights.read.bind(weights);
+  const a = 'memory_attention.';
+  const cells = grid * grid;
+  const linear = (x: string, p: string, out: number, into: number): string =>
+    node('linear', [x, read(`${p}weight`, [out, into]), 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 tokensOf = (map: string, channels: number): string =>
+    node('permute', [node('reshape', [map], { shape: [channels, cells] })], { order: [1, 0] });
+
+  const sine = sam2SinePositions(dim / 2, grid, grid);
+  const tenth = Float32Array.from(sine, (value) => Math.fround(Math.fround(0.1) * value));
+  const place = weights.constant(`sam2.feature_positions.${grid}`, [dim, grid, grid], tenth);
+  let x = tokensOf(node('add', [features, place]), dim);
+  const rotary = sam2Rotary(dim / config.attentionHeads, grid, grid);
+  if (config.attentionHeads !== 1)
+    throw new Error('the memory attention is written for one head, as SAM 2.1 has');
+  const cos = weights.constant(`sam2.rotary_cos.${grid}`, [cells * dim], rotary.cos);
+  const sin = weights.constant(`sam2.rotary_sin.${grid}`, [cells * dim], rotary.sin);
+  /* `x·cos + swap(x)·sin` over `count` blocks of the grid's cells, `swap` exchanging each pair. */
+  const turn = (value: string, count: number): string => {
+    const pairs = node('reshape', [value], { shape: [count * cells * (dim / 2), 2] });
+    const swapped = node(
+      'concat',
+      [
+        node('slice', [pairs], { axis: 1, start: 1, end: 2 }),
+        node('slice', [pairs], { axis: 1, start: 0, end: 1 }),
+      ],
+      { axis: 1 },
+    );
+    const flat = (v: string): string => node('reshape', [v], { shape: [count, cells * dim] });
+    const turned = node('add', [
+      node('mul', [flat(value), cos]),
+      node('mul', [flat(swapped), sin]),
+    ]);
+    return node('reshape', [turned], { shape: [count * cells, dim] });
+  };
+  const keyed = node('add', [memory, positions]);
+  const rotated = frames * cells;
+  const attend = (
+    p: string,
+    q: string,
+    k: string,
+    v: string,
+    into: number,
+    turnKeys: boolean,
+  ): string => {
+    const queries = turn(linear(q, `${p}q_proj.`, dim, dim), 1);
+    let keys = linear(k, `${p}k_proj.`, dim, into);
+    if (turnKeys) {
+      const cellsPart = turn(node('slice', [keys], { axis: 0, start: 0, end: rotated }), frames);
+      keys =
+        pointers === 0
+          ? cellsPart
+          : node(
+              'concat',
+              [
+                cellsPart,
+                node('slice', [keys], { axis: 0, start: rotated, end: rotated + pointers }),
+              ],
+              { axis: 0 },
+            );
+    } else {
+      keys = turn(keys, 1);
+    }
+    const values = linear(v, `${p}v_proj.`, dim, into);
+    return linear(
+      node('attention', [queries, keys, values], { heads: config.attentionHeads }),
+      `${p}o_proj.`,
+      dim,
+      dim,
+    );
+  };
+  for (let i = 0; i < config.attentionLayers; i += 1) {
+    const p = `${a}layers.${i}.`;
+    const q1 = norm(x, `${p}layer_norm1.`);
+    x = node('add', [x, attend(`${p}self_attn.`, q1, q1, q1, dim, false)]);
+    const q2 = norm(x, `${p}layer_norm2.`);
+    x = node('add', [x, attend(`${p}cross_attn_image.`, q2, keyed, memory, config.dim, true)]);
+    const q3 = norm(x, `${p}layer_norm3.`);
+    x = node('add', [
+      x,
+      linear(
+        node('relu', [linear(q3, `${p}linear1.`, config.feedForward, dim)]),
+        `${p}linear2.`,
+        dim,
+        config.feedForward,
+      ),
+    ]);
+  }
+  const normed = norm(x, `${a}layer_norm.`);
+  node(
+    'reshape',
+    [node('permute', [normed], { order: [1, 0] })],
+    { shape: [dim, grid, grid] },
+    output,
+  );
+}

package/src/models/sam2Positions.ts

@@ -0,0 +1,106 @@
+/**
+ * SAM 2's position tables, which depend only on shapes: the feature pyramid's and the memory's
+ * two-dimensional sines, the object pointers' one-dimensional sines, and the memory attention's
+ * rotary tables. Each is laid out once per size and used as a constant of a graph or by the host.
+ *
+ * **Every rounding is the upstream's, in single precision**: a coordinate normalised to the grid's
+ * extent, divided by its frequency, and only then taken through `sin` or `cos` — the engine's
+ * reproducible ones, since a capture reproduces. A frequency is `10000^t`, taken as `exp(t·ln 10000)`
+ * rather than `**`, which is not the same bits on every engine.
+ */
+import { exactCos, exactExp, exactSin } from '@driftengine/core';
+
+const f = Math.fround;
+const TWO_PI = f(2 * Math.PI);
+/** `ln 10000`, the temperature's logarithm. */
+const LOG_TEMPERATURE = 9.210340371976184;
+
+/* `10000^(2·⌊i/2⌋/count)` in single precision, as the upstream's `temperature ** (...)` makes it. */
+function frequency(i: number, count: number): number {
+  return f(exactExp(f((2 * Math.floor(i / 2)) / count) * LOG_TEMPERATURE));
+}
+
+/**
+ * `[2·features, height, width]`: the upstream's normalised sine embedding — a row's and a column's
+ * index from one, over the last and times 2π, each divided by `features` frequencies taken in pairs,
+ * sine then cosine alternating — with the rows' `features` channels before the columns'.
+ */
+export function sam2SinePositions(features: number, height: number, width: number): Float32Array {
+  const out = new Float32Array(2 * features * height * width);
+  const plane = height * width;
+  const embed = (index: number, count: number): number =>
+    f(f((index + 1) / f(count + 1e-6)) * TWO_PI);
+  for (let i = 0; i < features; i += 1) {
+    const divisor = frequency(i, features);
+    const wave = i % 2 === 0 ? exactSin : exactCos;
+    for (let y = 0; y < height; y += 1) {
+      const row = wave(f(embed(y, height) / divisor));
+      for (let x = 0; x < width; x += 1) {
+        out[i * plane + y * width + x] = row;
+        out[(features + i) * plane + y * width + x] = wave(f(embed(x, width) / divisor));
+      }
+    }
+  }
+  return out;
+}
+
+/**
+ * `[positions.length, dim]`: the upstream's `get_1d_sine_pe` — each position divided by `dim / 2`
+ * frequencies taken in pairs, its sines then its cosines.
+ */
+export function sam2PointerPositions(positions: readonly number[], dim: number): Float32Array {
+  const half = dim / 2;
+  const out = new Float32Array(positions.length * dim);
+  for (let p = 0; p < positions.length; p += 1) {
+    for (let k = 0; k < half; k += 1) {
+      const angle = f(f(positions[p] as number) / frequency(k, half));
+      out[p * dim + k] = exactSin(angle);
+      out[p * dim + half + k] = exactCos(angle);
+    }
+  }
+  return out;
+}
+
+/**
+ * The memory attention's axial rotary tables over a `width` by `height` grid, row by row, for a
+ * head of `dim`, at the upstream's base of 10,000 — its configuration's `rope_theta`, the one value
+ * any checkpoint here uses: `[cells, dim]` cosines and sines, the first half of the channels turning by the
+ * column and the second by the row, each frequency on a pair of channels. **The sine carries the
+ * rotation's sign** — negative on a pair's first channel, positive on its second — so a rotation is
+ * `x·cos + swap(x)·sin`, `swap` exchanging each pair, which a graph composes from its shape
+ * operators.
+ */
+export function sam2Rotary(
+  dim: number,
+  width: number,
+  height: number,
+): { readonly cos: Float32Array; readonly sin: Float32Array } {
+  const spatial = dim / 2;
+  const count = spatial / 2;
+  const inverse = Float32Array.from({ length: count }, (_, k) =>
+    f(1 / f(exactExp(f((2 * k) / spatial) * LOG_TEMPERATURE))),
+  );
+  const cells = width * height;
+  const cos = new Float32Array(cells * dim);
+  const sin = new Float32Array(cells * dim);
+  for (let y = 0; y < height; y += 1) {
+    for (let x = 0; x < width; x += 1) {
+      const at = (y * width + x) * dim;
+      for (let k = 0; k < count; k += 1) {
+        for (const [offset, position] of [
+          [0, x],
+          [spatial, y],
+        ] as const) {
+          const angle = f(position * (inverse[k] as number));
+          const c = exactCos(angle);
+          const s = exactSin(angle);
+          cos[at + offset + 2 * k] = c;
+          cos[at + offset + 2 * k + 1] = c;
+          sin[at + offset + 2 * k] = -s;
+          sin[at + offset + 2 * k + 1] = s;
+        }
+      }
+    }
+  }
+  return { cos, sin };
+}

package/src/models/sam2Tracker.ts

@@ -0,0 +1,300 @@
+/**
+ * SAM 2.1 tracking objects through a video: the upstream's `Sam2VideoModel` loop on the host, its
+ * networks as the engine's graphs, run by whatever the caller gives — the device's runner, or the
+ * CPU's evaluator.
+ *
+ * **A frame is decoded against what the tracker remembers.** A prompted frame, the first time, reads
+ * nothing and is given `no_memory_embedding`; every other frame's features are first conditioned by
+ * the memory attention on the frames and pointers `sam2Memories.ts` chooses, and `sam2Choice.ts`
+ * keeps one of the four masks decoded.
+ *
+ * **Then each object's frame becomes a memory**, all of a frame's objects together: the mask
+ * brought to the frame's size, binarised when any of them was just prompted and taken through a
+ * sigmoid otherwise — the upstream decides it once for the batch — encoded with the frame's
+ * features, the occlusion embedding added where the object is absent, and rounded to bfloat16.
+ *
+ * What it gives up: a mask as a prompt, which the upstream downsamples with an antialiased resize
+ * the runtime does not yet have, and more prompted frames than all of them — the upstream's default
+ * — are the two options not taken.
+ */
+import type { NetworkGraph, WeightSource } from '@driftengine/texture';
+import { graphFromWeights, type Architecture } from '@driftengine/texture';
+
+import type { GraphRun } from '../run.ts';
+
+import {
+  memoriesFor,
+  memoryTokens,
+  pointersFor,
+  toBfloat16,
+  type Sam2History,
+  type Sam2Output,
+} from './sam2Memories.ts';
+import { chooseMask } from './sam2Choice.ts';
+import { sam2SinePositions } from './sam2Positions.ts';
+import { sam21Decoder, sam21Encoder, type Sam21Config } from './sam21.ts';
+import { sam21MemoryAttention, sam21MemoryEncoder, sam21Upscale } from './sam21Video.ts';
+import { SAM2_VIDEO_PROMPT, samPromptTokens, samTokenCount, type SamPrompt } from './samPrompt.ts';
+
+export type { GraphRun } from '../run.ts';
+
+/** The weights each graph reads: a converted file's graphs, or a checkpoint for all four. */
+export interface Sam21Weights {
+  readonly encoder: WeightSource;
+  readonly decoder: WeightSource;
+  readonly memoryEncoder: WeightSource;
+  readonly memoryAttention: WeightSource;
+}
+
+export interface Sam21Video {
+  readonly frames: number;
+  /** The original frame's size, which prompts are given in. */
+  readonly height: number;
+  readonly width: number;
+  /** A frame prepared as the encoder takes it, `[3, size, size]`. */
+  frame(index: number): Float32Array;
+}
+
+/** A frame's result: every object's mask logits at a quarter of the encoder's size, and its score. */
+export interface Sam21Frame {
+  readonly frame: number;
+  readonly objects: readonly number[];
+  /** `[objects, (size / 4)²]`. */
+  readonly masks: Float32Array;
+  readonly scores: Float32Array;
+}
+
+const NO_WEIGHTS: WeightSource = { get: () => undefined, names: () => [] };
+
+interface Tracked extends Sam2History {
+  readonly prompts: Map<number, SamPrompt>;
+  readonly seen: Set<number>;
+}
+
+export class Sam21Tracker {
+  private readonly objects = new Map<number, Tracked>();
+  private fresh: number[] = [];
+  private readonly graphs = new Map<string, NetworkGraph>();
+  private encoded: {
+    readonly frame: number;
+    readonly values: ReadonlyMap<string, Float32Array>;
+  } | null = null;
+  private readonly cellPositions: Float32Array;
+
+  constructor(
+    private readonly config: Sam21Config,
+    private readonly weights: Sam21Weights,
+    private readonly run: GraphRun,
+    private readonly video: Sam21Video,
+  ) {
+    const grid = config.size / 16;
+    const width = config.memory.dim;
+    const sine = sam2SinePositions(width / 2, grid, grid);
+    this.cellPositions = new Float32Array(grid * grid * width);
+    for (let c = 0; c < width; c += 1) {
+      for (let i = 0; i < grid * grid; i += 1)
+        this.cellPositions[i * width + c] = sine[c * grid * grid + i] as number;
+    }
+  }
+
+  /** Prompts on `frame`, replacing each object's earlier ones there; the next step decodes them. */
+  prompt(frame: number, prompts: ReadonlyMap<number, SamPrompt>): void {
+    for (const [object, prompt] of prompts) {
+      let state = this.objects.get(object);
+      if (state === undefined) {
+        state = { prompts: new Map(), prompted: new Map(), tracked: new Map(), seen: new Set() };
+        this.objects.set(object, state);
+      }
+      state.prompts.set(frame, prompt);
+    }
+    this.fresh = [...prompts.keys()];
+  }
+
+  /** Every object on `frame`, as the upstream's `forward` over its session answers it. */
+  async step(frame: number, reverse = false): Promise<Sam21Frame> {
+    const objects = [...this.objects.keys()];
+    const side = this.config.size / 4;
+    const cells = side * side;
+    const masks = new Float32Array(objects.length * cells);
+    const scores = new Float32Array(objects.length);
+    const pending: { readonly output: Sam2Output; readonly fromPrompt: boolean }[] = [];
+    let backwards = reverse;
+    for (const [i, object] of objects.entries()) {
+      const state = this.objects.get(object) as Tracked;
+      const isNew = this.fresh.includes(object);
+      let output = state.prompted.get(frame);
+      let first = true;
+      if (isNew || output === undefined) {
+        first = false;
+        let prompt: SamPrompt | null = null;
+        if (isNew) {
+          first = !state.seen.has(frame);
+          if (first) backwards = false;
+          prompt = state.prompts.get(frame) ?? null;
+          if (prompt !== null) this.fresh = this.fresh.filter((id) => id !== object);
+        }
+        output = await this.decode(state, frame, first, prompt, backwards);
+        (first ? state.prompted : state.tracked).set(frame, output);
+        pending.push({ output, fromPrompt: prompt !== null });
+      }
+      masks.set(output.lowRes, i * cells);
+      scores[i] = output.score;
+      if (!first) state.seen.add(frame);
+    }
+    await this.remember(frame, pending);
+    return { frame, objects, masks, scores };
+  }
+
+  private graph(key: string, source: WeightSource, build: () => Architecture): NetworkGraph {
+    let graph = this.graphs.get(key);
+    if (graph === undefined) {
+      graph = graphFromWeights(source, build());
+      this.graphs.set(key, graph);
+    }
+    return graph;
+  }
+
+  private async features(frame: number): Promise<ReadonlyMap<string, Float32Array>> {
+    if (this.encoded?.frame !== frame) {
+      const graph = this.graph('encoder', this.weights.encoder, () => sam21Encoder(this.config));
+      const out = await this.run(graph, new Map([['image', this.video.frame(frame)]]));
+      this.encoded = {
+        frame,
+        values: new Map([...out].map(([name, values]) => [name, Float32Array.from(values)])),
+      };
+    }
+    return this.encoded.values;
+  }
+
+  private async decode(
+    state: Tracked,
+    frame: number,
+    first: boolean,
+    prompt: SamPrompt | null,
+    reverse: boolean,
+  ): Promise<Sam2Output> {
+    const { config } = this;
+    const { dim } = config.decoder;
+    const features = await this.features(frame);
+    let embedding = features.get('features') as Float32Array;
+    if (!first) embedding = await this.condition(state, frame, embedding, reverse);
+    const request = prompt ?? {};
+    const tokens = new Float32Array(samTokenCount(request, SAM2_VIDEO_PROMPT) * dim);
+    samPromptTokens(
+      this.weights.decoder,
+      dim,
+      config.size,
+      request,
+      this.video.height,
+      this.video.width,
+      tokens,
+      SAM2_VIDEO_PROMPT,
+    );
+    const n = tokens.length / dim;
+    const decoder = this.graph(`decoder.${n}.${first}`, this.weights.decoder, () =>
+      sam21Decoder(config, n, { noMemory: first }),
+    );
+    const out = await this.run(
+      decoder,
+      new Map([
+        ['embedding', embedding],
+        ['high1', features.get('high1') as Float32Array],
+        ['high0', features.get('high0') as Float32Array],
+        ['prompt', tokens],
+      ]),
+    );
+    const score = (out.get('object') as Float32Array)[0] as number;
+    const { lowRes, pointer } = chooseMask(
+      out.get('masks') as Float32Array,
+      out.get('quality') as Float32Array,
+      out.get('pointers') as Float32Array,
+      score,
+      (prompt?.points?.length ?? 0) + (prompt?.box === undefined ? 0 : 2),
+      this.tensor(this.weights.decoder, 'no_object_pointer'),
+    );
+    return { lowRes, pointer, score, memory: null };
+  }
+
+  private async condition(
+    state: Tracked,
+    frame: number,
+    features: Float32Array,
+    reverse: boolean,
+  ): Promise<Float32Array> {
+    const { config } = this;
+    const memories = memoriesFor(state, frame, config.memory.frames, reverse);
+    const pointers = pointersFor(state, frame, this.video.frames, config.memory.pointers, reverse);
+    const { memory, positions, pointerTokens } = memoryTokens(
+      this.weights.memoryAttention,
+      memories,
+      pointers,
+      this.cellPositions,
+      config.memory.dim,
+      config.decoder.dim,
+    );
+    const graph = this.graph(
+      `attention.${memories.length}.${pointerTokens}`,
+      this.weights.memoryAttention,
+      () => sam21MemoryAttention(config, memories.length, pointerTokens),
+    );
+    const out = await this.run(
+      graph,
+      new Map([
+        ['features', features],
+        ['memory', memory],
+        ['positions', positions],
+      ]),
+    );
+    return Float32Array.from(out.get('conditioned') as Float32Array);
+  }
+
+  private async remember(
+    frame: number,
+    pending: readonly { readonly output: Sam2Output; readonly fromPrompt: boolean }[],
+  ): Promise<void> {
+    if (pending.length === 0) return;
+    const { config } = this;
+    const binary = pending.some(({ fromPrompt }) => fromPrompt);
+    const features = (await this.features(frame)).get('features') as Float32Array;
+    const upscale = this.graph('upscale', NO_WEIGHTS, () => sam21Upscale(config, 1));
+    const encoder = this.graph(`memory.${binary}`, this.weights.memoryEncoder, () =>
+      sam21MemoryEncoder(config, binary),
+    );
+    const occlusion = this.tensor(
+      this.weights.memoryEncoder,
+      'occlusion_spatial_embedding_parameter',
+    );
+    const width = config.memory.dim;
+    for (const { output } of pending) {
+      const high = Float32Array.from(
+        (await this.run(upscale, new Map([['masks', output.lowRes]]))).get('high') as Float32Array,
+      );
+      if (binary)
+        for (let i = 0; i < high.length; i += 1) high[i] = (high[i] as number) > 0 ? 1 : 0;
+      const encoded = (
+        await this.run(
+          encoder,
+          new Map([
+            ['features', features],
+            ['mask', high],
+          ]),
+        )
+      ).get('memory') as Float32Array;
+      const cells = encoded.length / width;
+      const memory = new Float32Array(encoded.length);
+      for (let c = 0; c < width; c += 1) {
+        const add = output.score > 0 ? 0 : (occlusion[c] as number);
+        for (let i = 0; i < cells; i += 1)
+          memory[i * width + c] = Math.fround((encoded[c * cells + i] as number) + add);
+      }
+      toBfloat16(memory);
+      output.memory = memory;
+    }
+  }
+
+  private tensor(source: WeightSource, name: string): Float32Array {
+    const held = source.get(name);
+    if (held === undefined) throw new Error(`the tracker's weights have no "${name}"`);
+    return held.data;
+  }
+}