@driftengine/capture 4.0.0

package/src/models/owlv2Miniature.ts

@@ -0,0 +1,108 @@
+/**
+ * A seeded miniature of OWLv2: every tensor name of the upstream's detector, at a size the CPU
+ * evaluates in a moment — what the model's tests and its hand-run oracle share.
+ *
+ * **Small where the upstream lets it be**: a 64-pixel image in patches of 16, so a grid of four;
+ * towers two layers deep of 32 and 24 channels in two heads. **The text keeps
+ * the upstream's 16 positions**, since a query's length is what its causal mask and its pooling are
+ * about, and the vocabulary is three merges' worth — 517 tokens — which is also the tokenizer's
+ * hand-worked table (`clipTokenizer.test.ts`).
+ */
+import type { GraphTensor } from '@driftengine/texture';
+
+import type { Owlv2Config } from './owlv2.ts';
+import { draw } from './seeded.ts';
+
+export const MINIATURE_OWLV2: Owlv2Config = {
+  size: 64,
+  patch: 16,
+  vision: { dim: 32, heads: 2, depth: 2, hidden: 48 },
+  text: { dim: 24, heads: 2, depth: 2, hidden: 40, positions: 16 },
+};
+
+/** c+a, ca+t</w> and a+t</w>: tokens 512 to 514, and the start and end of text 515 and 516. */
+export const MINIATURE_OWLV2_MERGES = Float32Array.of(66, 64, 512, 339, 64, 339);
+
+/**
+ * Three queries as token ids: one ending early, one with the padding token inside it — the "!" a
+ * text may hold — and one filling every position.
+ */
+export const MINIATURE_OWLV2_QUERIES: readonly (readonly number[])[] = [
+  [515, 513, 516],
+  [515, 83, 514, 0, 320, 513, 516],
+  [515, 10, 200, 300, 400, 500, 511, 512, 1, 2, 3, 100, 250, 350, 450, 516],
+];
+
+/** The original image's size the boxes are scaled to: wider than tall, padded below. */
+export const MINIATURE_OWLV2_IMAGE = { height: 36, width: 48 } as const;
+
+/** Every tensor the upstream's `Owlv2ForObjectDetection` holds, named as its checkpoints are. */
+export function miniatureOwlv2Checkpoint(
+  config: Owlv2Config,
+  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 tower = (prefix: string, dim: number, hidden: number, depth: number): void => {
+    for (let i = 0; i < depth; i += 1) {
+      const p = `${prefix}encoder.layers.${i}.`;
+      for (const norm of ['layer_norm1', 'layer_norm2']) {
+        add(`${p}${norm}.weight`, [dim]);
+        add(`${p}${norm}.bias`, [dim]);
+      }
+      for (const projection of ['q_proj', 'k_proj', 'v_proj', 'out_proj']) {
+        add(`${p}self_attn.${projection}.weight`, [dim, dim]);
+        add(`${p}self_attn.${projection}.bias`, [dim]);
+      }
+      add(`${p}mlp.fc1.weight`, [hidden, dim]);
+      add(`${p}mlp.fc1.bias`, [hidden]);
+      add(`${p}mlp.fc2.weight`, [dim, hidden]);
+      add(`${p}mlp.fc2.bias`, [dim]);
+    }
+  };
+  const { vision, text, patch } = config;
+  const cells = (config.size / patch) * (config.size / patch);
+  const t = 'owlv2.text_model.';
+  add(`${t}embeddings.token_embedding.weight`, [
+    512 + MINIATURE_OWLV2_MERGES.length / 2 + 2,
+    text.dim,
+  ]);
+  add(`${t}embeddings.position_embedding.weight`, [text.positions, text.dim]);
+  tower(t, text.dim, text.hidden, text.depth);
+  add(`${t}final_layer_norm.weight`, [text.dim]);
+  add(`${t}final_layer_norm.bias`, [text.dim]);
+  const v = 'owlv2.vision_model.';
+  add(`${v}embeddings.class_embedding`, [vision.dim]);
+  add(`${v}embeddings.patch_embedding.weight`, [vision.dim, 3, patch, patch]);
+  add(`${v}embeddings.position_embedding.weight`, [cells + 1, vision.dim]);
+  for (const norm of ['pre_layernorm', 'post_layernorm']) {
+    add(`${v}${norm}.weight`, [vision.dim]);
+    add(`${v}${norm}.bias`, [vision.dim]);
+  }
+  tower(v, vision.dim, vision.hidden, vision.depth);
+  add('owlv2.visual_projection.weight', [text.dim, vision.dim]);
+  add('owlv2.text_projection.weight', [text.dim, text.dim]);
+  add('owlv2.logit_scale', []);
+  add('layer_norm.weight', [vision.dim]);
+  add('layer_norm.bias', [vision.dim]);
+  add('class_head.dense0.weight', [text.dim, vision.dim]);
+  add('class_head.dense0.bias', [text.dim]);
+  for (const logit of ['logit_shift', 'logit_scale']) {
+    add(`class_head.${logit}.weight`, [1, vision.dim]);
+    add(`class_head.${logit}.bias`, [1]);
+  }
+  for (const [head, last] of [
+    ['box_head', 4],
+    ['objectness_head', 1],
+  ] as const) {
+    add(`${head}.dense0.weight`, [vision.dim, vision.dim]);
+    add(`${head}.dense0.bias`, [vision.dim]);
+    add(`${head}.dense1.weight`, [vision.dim, vision.dim]);
+    add(`${head}.dense1.bias`, [vision.dim]);
+    add(`${head}.dense2.weight`, [last, vision.dim]);
+    add(`${head}.dense2.bias`, [last]);
+  }
+  return out;
+}

package/src/models/rope.ts

@@ -0,0 +1,111 @@
+/**
+ * Two-dimensional rotary position embedding, as Depth Anything 3's backbone applies it to queries
+ * and keys, composed from the runtime's operators.
+ *
+ * **The rotation is `x·cos + rotate(x)·sin`, per head, with a head's first half turned by the
+ * token's row and its second half by its column.** Within each half, `rotate` swaps the two
+ * quarters and negates the first — so with the sign folded into the sine table it is a swap, and a
+ * swap is two slices and a concatenation. The tables depend only on positions, which the graph's
+ * shapes fix, so they are constants of the graph, rebuilt at every size.
+ *
+ * **What it gives up is a dispatch count**: eight small copies a rotation where a kernel of its own
+ * would take one. At a depth model's shapes each is a few microseconds against an attention of a
+ * millisecond and a half, and every operator used here already has a device bound; a `rope`
+ * operator is the step to take when a measurement says the copies matter.
+ *
+ * **The tables are computed in single precision, operation by operation**, as the upstream's are:
+ * a frequency as `100^(k/16)` rounded, an angle as position times frequency rounded, then its
+ * cosine and sine. Computing them in double and rounding once would move an angle near 37 by
+ * two millionths, which a parity check would then have to allow for. The transcendentals are the
+ * engine's reproducible ones, so the tables are the same bits in every engine.
+ */
+import { exactCos, exactExp, exactSin } from '@driftengine/core';
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+
+const f = Math.fround;
+
+/*
+ * The base every rotary and sinusoidal embedding in these models is built on, as its logarithm:
+ * `100^e` is `exp(e · ln 100)`, through the engine's reproducible `exp`, since neither `**` nor
+ * `Math.pow` gives the same bits on every engine and a capture reproduces.
+ */
+export const LOG_BASE = 4.605170185988092;
+
+export interface RopeTables {
+  /** `[tokens, heads · headDim]`: the cosine of each channel's angle. */
+  readonly cos: Float32Array;
+  /** The same shape: the sine, negated where `rotate` negates. */
+  readonly sin: Float32Array;
+}
+
+/**
+ * The tables for tokens at `positions` — a row and a column each, as integers — over `heads` of
+ * `headDim` channels, at the upstream's base of 100.
+ */
+export function ropeTables(positions: Int32Array, heads: number, headDim: number): RopeTables {
+  const tokens = positions.length / 2;
+  const half = headDim / 2;
+  const quarter = half / 2;
+  const inverse = new Float32Array(quarter);
+  for (let k = 0; k < quarter; k += 1) {
+    inverse[k] = f(1 / f(exactExp(f((2 * k) / half) * LOG_BASE)));
+  }
+  const width = heads * headDim;
+  const cos = new Float32Array(tokens * width);
+  const sin = new Float32Array(tokens * width);
+  for (let t = 0; t < tokens; t += 1) {
+    for (let axis = 0; axis < 2; axis += 1) {
+      const position = positions[2 * t + axis] as number;
+      for (let j = 0; j < half; j += 1) {
+        const angle = f(position * (inverse[j % quarter] as number));
+        const c = f(exactCos(angle));
+        const s = f(exactSin(angle)) * (j < quarter ? -1 : 1);
+        for (let h = 0; h < heads; h += 1) {
+          const at = t * width + h * headDim + axis * half + j;
+          cos[at] = c;
+          sin[at] = s;
+        }
+      }
+    }
+  }
+  return { cos, sin };
+}
+
+/**
+ * The rotation of `x`, `[tokens, heads · headDim]`, by tables already added to the graph as the
+ * constants `cos` and `sin`.
+ */
+export function applyRope(
+  graph: GraphBuilder,
+  x: string,
+  tokens: number,
+  heads: number,
+  headDim: number,
+  cos: string,
+  sin: string,
+): string {
+  const quarter = headDim / 4;
+  const quarters = graph.node('reshape', [x], { shape: [tokens * heads * 2, 2, quarter] });
+  const first = graph.node('slice', [quarters], { axis: 1, start: 0, end: 1 });
+  const second = graph.node('slice', [quarters], { axis: 1, start: 1, end: 2 });
+  const swapped = graph.node('reshape', [graph.node('concat', [second, first], { axis: 1 })], {
+    shape: [tokens, heads * headDim],
+  });
+  return graph.node('add', [graph.node('mul', [x, cos]), graph.node('mul', [swapped, sin])]);
+}
+
+/** The tables as constants of the graph, named by `name`. */
+export function ropeConstants(
+  weights: Weights,
+  name: string,
+  positions: Int32Array,
+  heads: number,
+  headDim: number,
+): { readonly cos: string; readonly sin: string } {
+  const tables = ropeTables(positions, heads, headDim);
+  const shape = [positions.length / 2, heads * headDim];
+  return {
+    cos: weights.constant(`${name}.cos`, shape, tables.cos),
+    sin: weights.constant(`${name}.sin`, shape, tables.sin),
+  };
+}

package/src/models/sam21.ts

@@ -0,0 +1,182 @@
+/**
+ * SAM 2.1: masks from a point or a box, and masks carried through a video, as graphs — an encoder
+ * run once a frame, a decoder once a prompt, and the memory that carries an object from one frame
+ * to the next.
+ *
+ * **The encoder takes a frame at 1024 square**, resized without keeping its aspect, as the upstream's
+ * processor resizes it, and normalised by ImageNet's mean and deviation. It answers three levels of
+ * features — the decoder's input at a sixteenth of the frame and its two skips at an eighth and a
+ * quarter — which a video's memory then reads too.
+ */
+import type { Architecture } from '@driftengine/texture';
+
+import { hiera, type HieraConfig } from './hiera.ts';
+import { sam2Decoder } from './sam2Decoder.ts';
+import type { SamDecoderConfig } from './samDecoder.ts';
+
+export interface Sam21Config {
+  /** The side the encoder takes a frame at. */
+  readonly size: number;
+  readonly encoder: HieraConfig;
+  readonly decoder: SamDecoderConfig;
+  readonly memory: Sam21MemoryConfig;
+}
+
+export interface Sam21MemoryConfig {
+  /** A memory's channels: the encoder's are narrowed to these before they are kept. */
+  readonly dim: number;
+  readonly attentionLayers: number;
+  readonly attentionHeads: number;
+  readonly feedForward: number;
+  readonly fuserLayers: number;
+  readonly fuserHidden: number;
+  /** Frames of memory kept, the prompted frame among them, and object pointers read. */
+  readonly frames: number;
+  readonly pointers: number;
+}
+
+export const SAM_21_TINY: Sam21Config = {
+  size: 1024,
+  encoder: {
+    dims: [96, 192, 384, 768],
+    blocks: [1, 2, 7, 2],
+    heads: [1, 2, 4, 8],
+    windows: [8, 4, 14, 7],
+    global: [5, 7, 9],
+    background: 7,
+    mlpRatio: 4,
+    fpn: 256,
+    topDown: [2, 3],
+  },
+  decoder: {
+    dim: 256,
+    heads: 8,
+    mlp: 2048,
+    depth: 2,
+    downsample: 2,
+    masks: 4,
+    iouHidden: 256,
+    iouDepth: 3,
+    maskChannels: 16,
+  },
+  memory: {
+    dim: 64,
+    attentionLayers: 4,
+    attentionHeads: 1,
+    feedForward: 2048,
+    fuserLayers: 2,
+    fuserHidden: 1024,
+    frames: 7,
+    pointers: 16,
+  },
+};
+
+/** `image`, `[3, size, size]`, to `features`, `high1` and `high0`. */
+export function sam21Encoder(config: Sam21Config): Architecture {
+  return (weights, graph) => {
+    for (const prefix of [
+      'prompt_encoder.',
+      'shared_image_embedding.',
+      'memory_attention.',
+      'memory_encoder.',
+      'object_pointer_proj.',
+      'temporal_positional_encoding_projection_layer.',
+      'mask_downsample.',
+    ]) {
+      weights.ignore(prefix);
+    }
+    for (const name of [
+      'no_memory_embedding',
+      'no_memory_positional_encoding',
+      'memory_temporal_positional_encoding',
+      'no_object_pointer',
+      'occlusion_spatial_embedding_parameter',
+    ]) {
+      weights.ignore(name);
+    }
+    /* The decoder's, but for the two projections the encoder applies once a frame, which it reads. */
+    weights.ignore('mask_decoder.');
+    hiera(weights, graph, config.encoder, 'image', config.size, {
+      features: 'features',
+      high1: 'high1',
+      high0: 'high0',
+    });
+    return {
+      inputs: [{ name: 'image', shape: [3, config.size, config.size] }],
+      outputs: ['features', 'high1', 'high0'],
+    };
+  };
+}
+
+/** The prefixes and names of the memory's weights, which the image graphs set aside. */
+const MEMORY = [
+  'memory_attention.',
+  'memory_encoder.',
+  'mask_downsample.',
+  'temporal_positional_encoding_projection_layer.',
+  'memory_temporal_positional_encoding',
+  'no_memory_positional_encoding',
+  'occlusion_spatial_embedding_parameter',
+] as const;
+
+/**
+ * `embedding`, `high1`, `high0` and `prompt`, `[tokens, dim]` — absent when `tokens` is zero — and
+ * `mask` when `refine` is set, to `masks`, `quality`, `object` and `pointers`. With `noMemory` the
+ * decoder adds `no_memory_embedding` to the features itself, as an image with nothing remembered
+ * is decoded; otherwise the features are the memory attention's.
+ */
+export function sam21Decoder(
+  config: Sam21Config,
+  tokens: number,
+  options: { readonly refine?: boolean; readonly noMemory?: boolean } = {},
+): Architecture {
+  return (weights, graph) => {
+    weights.ignore('vision_encoder.');
+    weights.ignore('mask_decoder.conv_s0.');
+    weights.ignore('mask_decoder.conv_s1.');
+    for (const name of MEMORY) weights.ignore(name);
+    const { dim } = config.decoder;
+    const grid = config.size / 16;
+    /* Read whether this graph uses them or not: the host's pointer when no object is seen. */
+    weights.read('no_object_pointer', [1, dim]);
+    const noMemory = weights.read('no_memory_embedding', [1, 1, dim]);
+    let embedding = 'embedding';
+    if (options.noMemory === true) {
+      const flat = graph.node('reshape', [noMemory], { shape: [dim] });
+      const tokensFirst = graph.node(
+        'permute',
+        [graph.node('reshape', ['embedding'], { shape: [dim, grid * grid] })],
+        { order: [1, 0] },
+      );
+      embedding = graph.node(
+        'reshape',
+        [graph.node('permute', [graph.node('add', [tokensFirst, flat])], { order: [1, 0] })],
+        {
+          shape: [dim, grid, grid],
+        },
+      );
+    }
+    const inputs = [
+      { name: 'embedding', shape: [dim, grid, grid] },
+      { name: 'high1', shape: [dim / 4, 2 * grid, 2 * grid] },
+      { name: 'high0', shape: [dim / 8, 4 * grid, 4 * grid] },
+      ...(tokens > 0 ? [{ name: 'prompt', shape: [tokens, dim] }] : []),
+      ...(options.refine === true ? [{ name: 'mask', shape: [1, 4 * grid, 4 * grid] }] : []),
+    ];
+    sam2Decoder(
+      weights,
+      graph,
+      config.decoder,
+      grid,
+      {
+        embedding,
+        prompt: tokens > 0 ? 'prompt' : null,
+        mask: options.refine === true ? 'mask' : null,
+        high1: 'high1',
+        high0: 'high0',
+      },
+      { masks: 'masks', quality: 'quality', object: 'object', pointers: 'pointers' },
+    );
+    return { inputs, outputs: ['masks', 'quality', 'object', 'pointers'] };
+  };
+}

package/src/models/sam21Miniature.ts

@@ -0,0 +1,225 @@
+/**
+ * A seeded miniature of SAM 2.1: every tensor name and shape of the upstream's video model, at a
+ * size the CPU evaluates in a moment — what the model's tests and its hand-run oracle share.
+ *
+ * **Small where the upstream lets it be.** A 64-pixel frame, stages of 8, 16, 32 and 64 channels in
+ * 1, 2, 3 and 2 blocks, windows of 8, 4, 14 and 7 over grids of 16, 8, 4 and 2 — so windows divide,
+ * pad and outgrow their grid, one block attends globally and every stage change pools its queries.
+ * **The width is 128**, where the upstream's is 256: a memory is 64 channels wide, a width the
+ * upstream writes as a literal, and an object pointer is cut into width/64 tokens, so the width must
+ * be a multiple of 64 — and 128 cuts each pointer in two, as the real model's 256 cuts it in four.
+ */
+import type { GraphTensor } from '@driftengine/texture';
+
+import type { Sam21Config } from './sam21.ts';
+import type { SamPrompt } from './samPrompt.ts';
+import { draw } from './seeded.ts';
+
+export const MINIATURE_SAM_21: Sam21Config = {
+  size: 64,
+  encoder: {
+    dims: [8, 16, 32, 64],
+    blocks: [1, 2, 3, 2],
+    heads: [1, 2, 2, 4],
+    windows: [8, 4, 14, 7],
+    global: [4],
+    background: 7,
+    mlpRatio: 4,
+    fpn: 128,
+    topDown: [2, 3],
+  },
+  decoder: {
+    dim: 128,
+    heads: 2,
+    mlp: 32,
+    depth: 2,
+    downsample: 2,
+    masks: 4,
+    iouHidden: 16,
+    iouDepth: 3,
+    maskChannels: 16,
+  },
+  memory: {
+    dim: 64,
+    attentionLayers: 2,
+    attentionHeads: 1,
+    feedForward: 32,
+    fuserLayers: 2,
+    fuserHidden: 64,
+    frames: 7,
+    pointers: 16,
+  },
+};
+
+/**
+ * The original frame the prompts are given in, 36 by 48, which SAM 2 resizes to its square — so
+ * each axis scales by its own factor — and a point, a box, and both with a second point off the
+ * object.
+ */
+export const MINIATURE_SAM_21_IMAGE = { height: 36, width: 48 } as const;
+export const MINIATURE_SAM_21_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 `Sam2VideoModel` holds, named as its checkpoints are. */
+export function miniatureSam21Checkpoint(
+  config: Sam21Config,
+  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 linear = (p: string, outputs: number, inputs: number): void => {
+    add(`${p}weight`, [outputs, inputs]);
+    add(`${p}bias`, [outputs]);
+  };
+  const conv = (p: string, outputs: number, inputs: number, kernel: number): void => {
+    add(`${p}weight`, [outputs, inputs, kernel, kernel]);
+    add(`${p}bias`, [outputs]);
+  };
+  const norm = (p: string, dim: number): void => {
+    add(`${p}weight`, [dim]);
+    add(`${p}bias`, [dim]);
+  };
+  /* The upstream's two-to-n-layer perceptron: `proj_in`, then `layers`, then `proj_out`. */
+  const perceptron = (
+    p: string,
+    input: number,
+    hidden: number,
+    output: number,
+    depth: number,
+  ): void => {
+    linear(`${p}proj_in.`, hidden, input);
+    for (let i = 0; i < depth - 2; i += 1) linear(`${p}layers.${i}.`, hidden, hidden);
+    linear(`${p}proj_out.`, output, hidden);
+  };
+
+  const { encoder, decoder, memory } = config;
+  const width = decoder.dim;
+  add('no_memory_embedding', [1, 1, width]);
+  add('no_memory_positional_encoding', [1, 1, width]);
+  add('memory_temporal_positional_encoding', [memory.frames, 1, 1, memory.dim]);
+  add('no_object_pointer', [1, width]);
+  add('occlusion_spatial_embedding_parameter', [1, memory.dim]);
+  add('shared_image_embedding.positional_embedding', [2, width / 2]);
+
+  const b = 'vision_encoder.backbone.';
+  const d0 = encoder.dims[0];
+  add(`${b}pos_embed`, [1, d0, encoder.background, encoder.background]);
+  add(`${b}pos_embed_window`, [1, d0, encoder.windows[0], encoder.windows[0]]);
+  conv(`${b}patch_embed.projection.`, d0, 3, 7);
+  let index = 0;
+  for (let stage = 0; stage < 4; stage += 1) {
+    for (let block = 0; block < (encoder.blocks[stage] as number); block += 1) {
+      const p = `${b}blocks.${index}.`;
+      const first = stage > 0 && block === 0;
+      const dim = first ? (encoder.dims[stage - 1] as number) : (encoder.dims[stage] as number);
+      const outDim = encoder.dims[stage] as number;
+      norm(`${p}layer_norm1.`, dim);
+      linear(`${p}attn.qkv.`, 3 * outDim, dim);
+      linear(`${p}attn.proj.`, outDim, outDim);
+      norm(`${p}layer_norm2.`, outDim);
+      linear(`${p}mlp.proj_in.`, outDim * encoder.mlpRatio, outDim);
+      linear(`${p}mlp.proj_out.`, outDim, outDim * encoder.mlpRatio);
+      if (first) linear(`${p}proj.`, outDim, dim);
+      index += 1;
+    }
+  }
+  for (let i = 0; i < 4; i += 1)
+    conv(`vision_encoder.neck.convs.${i}.`, encoder.fpn, encoder.dims[3 - i] as number, 1);
+
+  const pe = 'prompt_encoder.';
+  add(`${pe}shared_embedding.positional_embedding`, [2, width / 2]);
+  const c = decoder.maskChannels;
+  conv(`${pe}mask_embed.conv1.`, c / 4, 1, 2);
+  conv(`${pe}mask_embed.conv2.`, c, c / 4, 2);
+  conv(`${pe}mask_embed.conv3.`, width, c, 1);
+  norm(`${pe}mask_embed.layer_norm1.`, c / 4);
+  norm(`${pe}mask_embed.layer_norm2.`, c);
+  add(`${pe}no_mask_embed.weight`, [1, width]);
+  add(`${pe}point_embed.weight`, [4, width]);
+  add(`${pe}not_a_point_embed.weight`, [1, width]);
+
+  const m = 'mask_decoder.';
+  const t = `${m}transformer.`;
+  const inner = width / decoder.downsample;
+  const attention = (p: string, dim: number, kv = width): void => {
+    linear(`${p}q_proj.`, dim, width);
+    linear(`${p}k_proj.`, dim, kv);
+    linear(`${p}v_proj.`, dim, kv);
+    linear(`${p}o_proj.`, width, dim);
+  };
+  add(`${m}iou_token.weight`, [1, width]);
+  add(`${m}mask_tokens.weight`, [decoder.masks, width]);
+  for (let i = 0; i < decoder.depth; i += 1) {
+    const p = `${t}layers.${i}.`;
+    attention(`${p}self_attn.`, width);
+    norm(`${p}layer_norm1.`, width);
+    attention(`${p}cross_attn_token_to_image.`, inner);
+    norm(`${p}layer_norm2.`, width);
+    perceptron(`${p}mlp.`, width, decoder.mlp, width, decoder.depth);
+    norm(`${p}layer_norm3.`, width);
+    norm(`${p}layer_norm4.`, width);
+    attention(`${p}cross_attn_image_to_token.`, inner);
+  }
+  attention(`${t}final_attn_token_to_image.`, inner);
+  norm(`${t}layer_norm_final_attn.`, width);
+  add(`${m}upscale_conv1.weight`, [width, width / 4, 2, 2]);
+  add(`${m}upscale_conv1.bias`, [width / 4]);
+  add(`${m}upscale_conv2.weight`, [width / 4, width / 8, 2, 2]);
+  add(`${m}upscale_conv2.bias`, [width / 8]);
+  norm(`${m}upscale_layer_norm.`, width / 4);
+  for (let i = 0; i < decoder.masks; i += 1) {
+    perceptron(`${m}output_hypernetworks_mlps.${i}.`, width, width, width / 8, 3);
+  }
+  perceptron(`${m}iou_prediction_head.`, width, decoder.iouHidden, decoder.masks, decoder.iouDepth);
+  conv(`${m}conv_s0.`, width / 8, width, 1);
+  conv(`${m}conv_s1.`, width / 4, width, 1);
+  add(`${m}obj_score_token.weight`, [1, width]);
+  perceptron(`${m}pred_obj_score_head.`, width, width, 1, 3);
+
+  const a = 'memory_attention.';
+  for (let i = 0; i < memory.attentionLayers; i += 1) {
+    const p = `${a}layers.${i}.`;
+    attention(`${p}self_attn.`, width);
+    attention(`${p}cross_attn_image.`, width, memory.dim);
+    linear(`${p}linear1.`, memory.feedForward, width);
+    linear(`${p}linear2.`, width, memory.feedForward);
+    norm(`${p}layer_norm1.`, width);
+    norm(`${p}layer_norm2.`, width);
+    norm(`${p}layer_norm3.`, width);
+  }
+  norm(`${a}layer_norm.`, width);
+
+  const e = 'memory_encoder.';
+  let channels = 1;
+  for (let i = 0; i < 4; i += 1) {
+    conv(`${e}mask_downsampler.layers.${i}.conv.`, channels * 4, channels, 3);
+    norm(`${e}mask_downsampler.layers.${i}.layer_norm.`, channels * 4);
+    channels *= 4;
+  }
+  conv(`${e}mask_downsampler.final_conv.`, width, channels, 1);
+  conv(`${e}feature_projection.`, width, width, 1);
+  for (let i = 0; i < memory.fuserLayers; i += 1) {
+    const p = `${e}memory_fuser.layers.${i}.`;
+    add(`${p}scale`, [width]);
+    conv(`${p}depthwise_conv.`, width, 1, 7);
+    norm(`${p}layer_norm.`, width);
+    linear(`${p}pointwise_conv1.`, memory.fuserHidden, width);
+    linear(`${p}pointwise_conv2.`, width, memory.fuserHidden);
+  }
+  conv(`${e}projection.`, memory.dim, width, 1);
+  conv('mask_downsample.', 1, 1, 4);
+  perceptron('object_pointer_proj.', width, width, width, 3);
+  linear('temporal_positional_encoding_projection_layer.', memory.dim, width);
+  return out;
+}