@driftengine/capture 4.0.0

package/dist/models/owlv2Miniature.js

@@ -0,0 +1,87 @@
+import { draw } from './seeded.js';
+export const MINIATURE_OWLV2 = {
+    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 = [
+    [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 };
+/** Every tensor the upstream's `Owlv2ForObjectDetection` holds, named as its checkpoints are. */
+export function miniatureOwlv2Checkpoint(config, seed) {
+    const out = new Map();
+    const add = (name, shape) => {
+        out.set(name, { shape, data: draw(name, shape, seed) });
+    };
+    const tower = (prefix, dim, hidden, depth) => {
+        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],
+    ]) {
+        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/dist/models/rope.d.ts

@@ -0,0 +1,23 @@
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+export declare 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 declare function ropeTables(positions: Int32Array, heads: number, headDim: number): RopeTables;
+/**
+ * The rotation of `x`, `[tokens, heads · headDim]`, by tables already added to the graph as the
+ * constants `cos` and `sin`.
+ */
+export declare function applyRope(graph: GraphBuilder, x: string, tokens: number, heads: number, headDim: number, cos: string, sin: string): string;
+/** The tables as constants of the graph, named by `name`. */
+export declare function ropeConstants(weights: Weights, name: string, positions: Int32Array, heads: number, headDim: number): {
+    readonly cos: string;
+    readonly sin: string;
+};

package/dist/models/rope.js

@@ -0,0 +1,84 @@
+/**
+ * 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';
+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;
+/**
+ * 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, heads, headDim) {
+    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];
+            for (let j = 0; j < half; j += 1) {
+                const angle = f(position * inverse[j % quarter]);
+                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, x, tokens, heads, headDim, cos, sin) {
+    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, name, positions, heads, headDim) {
+    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/dist/models/sam21.d.ts

@@ -0,0 +1,45 @@
+/**
+ * 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 { type HieraConfig } from './hiera.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 declare const SAM_21_TINY: Sam21Config;
+/** `image`, `[3, size, size]`, to `features`, `high1` and `high0`. */
+export declare function sam21Encoder(config: Sam21Config): Architecture;
+/**
+ * `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 declare function sam21Decoder(config: Sam21Config, tokens: number, options?: {
+    readonly refine?: boolean;
+    readonly noMemory?: boolean;
+}): Architecture;

package/dist/models/sam21.js

@@ -0,0 +1,126 @@
+import { hiera } from './hiera.js';
+import { sam2Decoder } from './sam2Decoder.js';
+export const SAM_21_TINY = {
+    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) {
+    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',
+];
+/**
+ * `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, tokens, options = {}) {
+    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/dist/models/sam21Miniature.d.ts

@@ -0,0 +1,37 @@
+/**
+ * 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';
+export declare const MINIATURE_SAM_21: Sam21Config;
+/**
+ * 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 declare const MINIATURE_SAM_21_IMAGE: {
+    readonly height: 36;
+    readonly width: 48;
+};
+export declare const MINIATURE_SAM_21_PROMPTS: {
+    readonly point: {
+        readonly points: readonly [readonly [20.5, 14.25, 1]];
+    };
+    readonly box: {
+        readonly box: readonly [6, 5, 40.5, 30];
+    };
+    readonly both: {
+        readonly points: readonly [readonly [20.5, 14.25, 1], readonly [33, 25.75, 0]];
+        readonly box: readonly [6, 5, 40.5, 30];
+    };
+};
+/** Every tensor the upstream's `Sam2VideoModel` holds, named as its checkpoints are. */
+export declare function miniatureSam21Checkpoint(config: Sam21Config, seed: number): Map<string, GraphTensor>;

package/dist/models/sam21Miniature.js

@@ -0,0 +1,194 @@
+import { draw } from './seeded.js';
+export const MINIATURE_SAM_21 = {
+    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 };
+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],
+    },
+};
+/** Every tensor the upstream's `Sam2VideoModel` holds, named as its checkpoints are. */
+export function miniatureSam21Checkpoint(config, seed) {
+    const out = new Map();
+    const add = (name, shape) => {
+        out.set(name, { shape, data: draw(name, shape, seed) });
+    };
+    const linear = (p, outputs, inputs) => {
+        add(`${p}weight`, [outputs, inputs]);
+        add(`${p}bias`, [outputs]);
+    };
+    const conv = (p, outputs, inputs, kernel) => {
+        add(`${p}weight`, [outputs, inputs, kernel, kernel]);
+        add(`${p}bias`, [outputs]);
+    };
+    const norm = (p, dim) => {
+        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, input, hidden, output, depth) => {
+        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]; block += 1) {
+            const p = `${b}blocks.${index}.`;
+            const first = stage > 0 && block === 0;
+            const dim = first ? encoder.dims[stage - 1] : encoder.dims[stage];
+            const outDim = encoder.dims[stage];
+            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], 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, dim, kv = width) => {
+        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;
+}

package/dist/models/sam21Video.d.ts

@@ -0,0 +1,30 @@
+/**
+ * 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';
+/**
+ * `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 declare function sam21MemoryEncoder(config: Sam21Config, binary: boolean): Architecture;
+/**
+ * `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 declare function sam21MemoryAttention(config: Sam21Config, frames: number, pointers: number): Architecture;
+/** `masks`, `[count, 4·grid, 4·grid]`, to `high`, `[count, size, size]`, bilinearly. Reads no weights. */
+export declare function sam21Upscale(config: Sam21Config, count: number): Architecture;
+/**
+ * `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 declare function sam21MasksToImage(config: Sam21Config, count: number, height: number, width: number): Architecture;