@driftengine/capture 4.0.0

package/dist/models/samMiniature.js

@@ -0,0 +1,161 @@
+import { draw } from './seeded.js';
+export const MINIATURE_MOBILE_SAM = {
+    size: 64,
+    encoder: {
+        embedDims: [8, 16, 24, 320],
+        depths: [1, 2, 1, 1],
+        heads: [1, 2, 3, 10],
+        windows: [7, 3, 4, 2],
+        mergeStrides: [2, 2, 1],
+        mlpRatio: 4,
+        expandRatio: 4,
+        neck: 256,
+    },
+    decoder: {
+        dim: 256,
+        heads: 8,
+        mlp: 32,
+        depth: 2,
+        downsample: 2,
+        masks: 4,
+        iouHidden: 16,
+        iouDepth: 3,
+        maskChannels: 16,
+    },
+};
+/** The classes of the miniature's classification head, which nothing runs. */
+export const MINIATURE_SAM_CLASSES = 10;
+/**
+ * The original image the prompts are given in: 36 by 48, prepared at 48 by 64 in the encoder's 64
+ * square, so the masks are cropped before they come down to it and coordinates are scaled by 4/3.
+ */
+export const MINIATURE_SAM_IMAGE = { height: 36, width: 48 };
+/** A point, which a padding point follows; a box; and two points with a box, which none follows. */
+export const MINIATURE_SAM_PROMPTS = {
+    point: { points: [[20.5, 14.25, 1]] },
+    box: { box: [6, 5, 40.5, 30] },
+    both: {
+        points: [
+            [20.5, 14.25, 1],
+            [33, 25.75, 0],
+        ],
+        box: [6, 5, 40.5, 30],
+    },
+};
+/** Every tensor the upstream's encoder, prompt encoder and mask decoder hold. */
+export function miniatureSamCheckpoint(config, seed) {
+    const out = new Map();
+    const add = (name, shape) => {
+        out.set(name, { shape, data: draw(name, shape, seed) });
+    };
+    const convBn = (p, cout, cin, kernel) => {
+        add(`${p}c.weight`, [cout, cin, kernel, kernel]);
+        for (const part of ['weight', 'bias', 'running_mean', 'running_var'])
+            add(`${p}bn.${part}`, [cout]);
+        out.set(`${p}bn.num_batches_tracked`, { shape: [], data: new Float32Array(1) });
+    };
+    const linear = (p, outputs, inputs) => {
+        add(`${p}weight`, [outputs, inputs]);
+        add(`${p}bias`, [outputs]);
+    };
+    const norm = (p, dim) => {
+        add(`${p}weight`, [dim]);
+        add(`${p}bias`, [dim]);
+    };
+    const { embedDims: dims, depths, windows, mlpRatio, expandRatio, neck } = config.encoder;
+    const e = 'image_encoder.';
+    convBn(`${e}patch_embed.seq.0.`, dims[0] / 2, 3, 3);
+    convBn(`${e}patch_embed.seq.2.`, dims[0], dims[0] / 2, 3);
+    for (let i = 0; i < depths[0]; i += 1) {
+        const p = `${e}layers.0.blocks.${i}.`;
+        const hidden = dims[0] * expandRatio;
+        convBn(`${p}conv1.`, hidden, dims[0], 1);
+        convBn(`${p}conv2.`, hidden, 1, 3);
+        convBn(`${p}conv3.`, dims[0], hidden, 1);
+    }
+    for (let stage = 0; stage < 4; stage += 1) {
+        const dim = dims[stage];
+        if (stage > 0) {
+            for (let i = 0; i < depths[stage]; i += 1) {
+                const p = `${e}layers.${stage}.blocks.${i}.`;
+                const window = windows[stage];
+                add(`${p}attn.attention_biases`, [config.encoder.heads[stage], window * window]);
+                norm(`${p}attn.norm.`, dim);
+                linear(`${p}attn.qkv.`, 3 * dim, dim);
+                linear(`${p}attn.proj.`, dim, dim);
+                norm(`${p}mlp.norm.`, dim);
+                linear(`${p}mlp.fc1.`, dim * mlpRatio, dim);
+                linear(`${p}mlp.fc2.`, dim, dim * mlpRatio);
+                convBn(`${p}local_conv.`, dim, 1, 3);
+            }
+        }
+        if (stage < 3) {
+            const next = dims[stage + 1];
+            const p = `${e}layers.${stage}.downsample.`;
+            convBn(`${p}conv1.`, next, dim, 1);
+            convBn(`${p}conv2.`, next, 1, 3);
+            convBn(`${p}conv3.`, next, next, 1);
+        }
+    }
+    norm(`${e}norm_head.`, dims[3]);
+    linear(`${e}head.`, MINIATURE_SAM_CLASSES, dims[3]);
+    add(`${e}neck.0.weight`, [neck, dims[3], 1, 1]);
+    norm(`${e}neck.1.`, neck);
+    add(`${e}neck.2.weight`, [neck, neck, 3, 3]);
+    norm(`${e}neck.3.`, neck);
+    const { dim, masks, mlp, maskChannels: c, iouHidden, iouDepth } = config.decoder;
+    const pe = 'prompt_encoder.';
+    add(`${pe}pe_layer.positional_encoding_gaussian_matrix`, [2, dim / 2]);
+    for (let i = 0; i < 4; i += 1)
+        add(`${pe}point_embeddings.${i}.weight`, [1, dim]);
+    add(`${pe}not_a_point_embed.weight`, [1, dim]);
+    add(`${pe}mask_downscaling.0.weight`, [c / 4, 1, 2, 2]);
+    add(`${pe}mask_downscaling.0.bias`, [c / 4]);
+    norm(`${pe}mask_downscaling.1.`, c / 4);
+    add(`${pe}mask_downscaling.3.weight`, [c, c / 4, 2, 2]);
+    add(`${pe}mask_downscaling.3.bias`, [c]);
+    norm(`${pe}mask_downscaling.4.`, c);
+    add(`${pe}mask_downscaling.6.weight`, [dim, c, 1, 1]);
+    add(`${pe}mask_downscaling.6.bias`, [dim]);
+    add(`${pe}no_mask_embed.weight`, [1, dim]);
+    const d = 'mask_decoder.';
+    const t = `${d}transformer.`;
+    const attention = (p, inner) => {
+        for (const projection of ['q_proj', 'k_proj', 'v_proj'])
+            linear(`${p}${projection}.`, inner, dim);
+        linear(`${p}out_proj.`, dim, inner);
+    };
+    const inner = dim / config.decoder.downsample;
+    for (let i = 0; i < config.decoder.depth; i += 1) {
+        const p = `${t}layers.${i}.`;
+        attention(`${p}self_attn.`, dim);
+        norm(`${p}norm1.`, dim);
+        attention(`${p}cross_attn_token_to_image.`, inner);
+        norm(`${p}norm2.`, dim);
+        linear(`${p}mlp.lin1.`, mlp, dim);
+        linear(`${p}mlp.lin2.`, dim, mlp);
+        norm(`${p}norm3.`, dim);
+        norm(`${p}norm4.`, dim);
+        attention(`${p}cross_attn_image_to_token.`, inner);
+    }
+    attention(`${t}final_attn_token_to_image.`, inner);
+    norm(`${t}norm_final_attn.`, dim);
+    add(`${d}iou_token.weight`, [1, dim]);
+    add(`${d}mask_tokens.weight`, [masks, dim]);
+    add(`${d}output_upscaling.0.weight`, [dim, dim / 4, 2, 2]);
+    add(`${d}output_upscaling.0.bias`, [dim / 4]);
+    norm(`${d}output_upscaling.1.`, dim / 4);
+    add(`${d}output_upscaling.3.weight`, [dim / 4, dim / 8, 2, 2]);
+    add(`${d}output_upscaling.3.bias`, [dim / 8]);
+    for (let i = 0; i < masks; i += 1) {
+        const p = `${d}output_hypernetworks_mlps.${i}.layers.`;
+        linear(`${p}0.`, dim, dim);
+        linear(`${p}1.`, dim, dim);
+        linear(`${p}2.`, dim / 8, dim);
+    }
+    const sizes = [dim, ...Array.from({ length: iouDepth - 1 }, () => iouHidden), masks];
+    for (let i = 0; i + 1 < sizes.length; i += 1) {
+        linear(`${d}iou_prediction_head.layers.${i}.`, sizes[i + 1], sizes[i]);
+    }
+    return out;
+}

package/dist/models/samPrompt.d.ts

@@ -0,0 +1,57 @@
+import type { WeightSource } from '@driftengine/texture';
+export interface SamPrompt {
+    /** Points in the original image's pixels, each labelled 1 on the object and 0 off it. */
+    readonly points?: readonly (readonly [x: number, y: number, label: 0 | 1])[];
+    /** A box in the original image's pixels: left, top, right, bottom. */
+    readonly box?: readonly [number, number, number, number];
+}
+/**
+ * Where a checkpoint keeps its prompt encoder, and how its processor places a prompt: SAM's and
+ * MobileSAM's, or Transformers' SAM 2.
+ */
+export interface SamPromptLayout {
+    /** The Fourier matrix, `[2, dim / 2]`. */
+    readonly gaussian: string;
+    /** A learned point embedding by kind — 0 off, 1 on, 2 and 3 a box's corners — as name and row. */
+    readonly embedding: (kind: number) => readonly [name: string, row: number];
+    readonly notAPoint: string;
+    /** Whether a box is followed by a padding point, as SAM 2's is. */
+    readonly boxPadding: boolean;
+    /**
+     * Whether a box's corners come first as points and a padding point always ends the prompt, as a
+     * SAM 2 video session orders them — where a prompt of nothing is a point meaning none, padded.
+     */
+    readonly boxAsPoints: boolean;
+    /**
+     * The prepared image's size for an original one: the longest side to `size` for SAM, the whole
+     * frame to a `size` square for SAM 2 — and whether a coordinate is scaled in double precision, as
+     * SAM's predictor does with numpy, or single, as SAM 2's processor does with a tensor.
+     */
+    readonly prepared: (height: number, width: number, size: number) => readonly [number, number];
+    readonly singleScale: boolean;
+}
+export declare const SAM_PROMPT: SamPromptLayout;
+export declare const SAM2_PROMPT: SamPromptLayout;
+/** SAM 2 in a video session: the same encoder, a box as its two corner points first. */
+export declare const SAM2_VIDEO_PROMPT: SamPromptLayout;
+/** The upstream's `get_preprocess_shape`: the longest side to `longest`, each side rounded. */
+export declare function samInputSize(height: number, width: number, longest: number): readonly [number, number];
+/** How many tokens a prompt is, in a layout. */
+export declare function samTokenCount(request: SamPrompt, layout?: SamPromptLayout): number;
+/**
+ * `[grid · grid, dim]`: the position of each cell of the embedding's grid, at its centre, as the
+ * upstream's `get_dense_pe` lays it out and the decoder reads it as tokens.
+ */
+export declare function samGridPositions(gaussian: Float32Array, dim: number, grid: number): Float32Array;
+/**
+ * The tokens of `request`, `[samTokenCount(request, layout), dim]`, into `out`, for an original
+ * image of `height` by `width` prepared at `size` as `layout`'s processor prepares it.
+ */
+export declare function samPromptTokens(weights: WeightSource, dim: number, size: number, request: SamPrompt, height: number, width: number, out: Float32Array, layout?: SamPromptLayout): void;
+/**
+ * The upstream's stability score: the share of the pixels above `threshold − offset` that are also
+ * above `threshold + offset` — one over the other of the mask thresholded high and thresholded low,
+ * since one always lies inside the other. A mask with no pixel above the lower threshold scores
+ * NaN, as the upstream's does.
+ */
+export declare function stabilityScore(logits: Float32Array, threshold: number, offset: number): number;

package/dist/models/samPrompt.js

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

package/dist/models/seeded.d.ts

@@ -0,0 +1,16 @@
+export declare function nameSeed(name: string, seed: number): number;
+/**
+ * The spread a tensor's values are drawn with, by what the name says it is: a norm's scale near
+ * one, a layer scale near a half, a bias small, and a weight at 1.7/√fan-in. At 1/√fan-in the
+ * head's twenty convolutions shrank the signal until the depth was its biases — 1.08 to 1.11
+ * everywhere, and a slip in a norm's ε moved nothing a test could see; at He's √(6/fan-in) the
+ * fusion's residual sums compounded it to depths of 10^33. At 1.7 the depth spans 0.96 to 11.7.
+ *
+ * **A one-dimensional weight is a norm's scale** in every checkpoint here, whatever its name — a
+ * channel norm is often a bare index in a sequence — and is drawn near one.
+ *
+ * **A batch norm's statistics are drawn as statistics**: a running variance between a half and one
+ * and a half, so the scale folded from it stays near one, and a running mean small. Drawn as any
+ * other vector, a variance near zero would fold into a scale of hundreds.
+ */
+export declare function draw(name: string, shape: readonly number[], seed: number): Float32Array;

package/dist/models/seeded.js

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

package/dist/models/tinyVit.d.ts

@@ -0,0 +1,38 @@
+/**
+ * TinyViT as MobileSAM runs it: an image in, a 256-channel embedding at a sixteenth of its size out.
+ *
+ * **Four stages, the first convolutional.** A patch embedding of two strided convolutions takes the
+ * image to a quarter; the first stage is inverted residual blocks, each a pointwise expansion, a
+ * depthwise convolution and a pointwise projection; the other three are transformer blocks that
+ * attend within windows, each followed by a depthwise convolution over the grid and an MLP. A merge
+ * between stages widens the channels and halves the grid — except where the upstream keeps its
+ * stride at one, which it decides by the output width and this definition takes as `mergeStrides`.
+ *
+ * **Every convolution carries a batch norm**, folded in at conversion by `convBatchNorm`, and the
+ * transformer blocks attend within windows through `windowAttention`, which pads, partitions and
+ * biases them as the upstream does. The blocks' MLP norms take PyTorch's default ε, 1e-5, and the
+ * neck's channel norms 1e-6.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+type Four = readonly [number, number, number, number];
+export interface TinyVitConfig {
+    readonly embedDims: Four;
+    readonly depths: Four;
+    /** Heads per stage; the first stage is convolutional and has none. */
+    readonly heads: Four;
+    readonly windows: Four;
+    /** The stride of each merge between stages. */
+    readonly mergeStrides: readonly [number, number, number];
+    readonly mlpRatio: number;
+    readonly expandRatio: number;
+    /** The embedding's channels. */
+    readonly neck: number;
+}
+/** The grid the encoder's embedding has for an image `size` across. */
+export declare function tinyVitGrid(config: TinyVitConfig, size: number): number;
+/**
+ * `[neck, grid, grid]` from `image`, `[3, size, size]`, normalised and padded as the upstream
+ * prepares it; the embedding is written as `output`.
+ */
+export declare function tinyVit(weights: Weights, graph: GraphBuilder, config: TinyVitConfig, prefix: string, image: string, size: number, output: string): void;
+export {};

package/dist/models/tinyVit.js

@@ -0,0 +1,82 @@
+import { convBatchNorm } from './batchNorm.js';
+import { windowAttention } from './windowAttention.js';
+const EPSILON = 1e-5;
+const CHANNEL_EPSILON = 1e-6;
+/** The grid the encoder's embedding has for an image `size` across. */
+export function tinyVitGrid(config, size) {
+    return config.mergeStrides.reduce((grid, stride) => grid / stride, size / 4);
+}
+/**
+ * `[neck, grid, grid]` from `image`, `[3, size, size]`, normalised and padded as the upstream
+ * prepares it; the embedding is written as `output`.
+ */
+export function tinyVit(weights, graph, config, prefix, image, size, output) {
+    const node = graph.node.bind(graph);
+    const dims = config.embedDims;
+    weights.ignore(`${prefix}norm_head.`);
+    weights.ignore(`${prefix}head.`);
+    const convBn = (p, x, cin, cout, kernel, stride, groups) => convBatchNorm(weights, graph, p, x, cin, cout, kernel, stride, groups);
+    const gelu = (x) => node('gelu', [x]);
+    const tokensOf = (map, channels, cells) => node('permute', [node('reshape', [map], { shape: [channels, cells] })], { order: [1, 0] });
+    const mapOf = (tokens, channels, rows, cols) => node('reshape', [node('permute', [tokens], { order: [1, 0] })], {
+        shape: [channels, rows, cols],
+    });
+    const norm = (x, p, dim, epsilon) => node('layerNorm', [x, weights.read(`${p}weight`, [dim]), weights.read(`${p}bias`, [dim])], {
+        epsilon,
+    });
+    const linear = (x, p, out, inputs) => node('linear', [x, weights.read(`${p}weight`, [out, inputs]), weights.read(`${p}bias`, [out])]);
+    const merge = (p, map, dim, out, stride) => {
+        const expanded = gelu(convBn(`${p}conv1.`, map, dim, out, 1, 1, 1));
+        const strided = gelu(convBn(`${p}conv2.`, expanded, out, out, 3, stride, out));
+        return convBn(`${p}conv3.`, strided, out, out, 1, 1, 1);
+    };
+    let map = gelu(convBn(`${prefix}patch_embed.seq.0.`, image, 3, dims[0] / 2, 3, 2, 1));
+    map = convBn(`${prefix}patch_embed.seq.2.`, map, dims[0] / 2, dims[0], 3, 2, 1);
+    let grid = size / 4;
+    const hidden = dims[0] * config.expandRatio;
+    for (let i = 0; i < config.depths[0]; i += 1) {
+        const p = `${prefix}layers.0.blocks.${i}.`;
+        const inner = gelu(convBn(`${p}conv1.`, map, dims[0], hidden, 1, 1, 1));
+        const spread = gelu(convBn(`${p}conv2.`, inner, hidden, hidden, 3, 1, hidden));
+        map = gelu(node('add', [convBn(`${p}conv3.`, spread, hidden, dims[0], 1, 1, 1), map]));
+    }
+    map = merge(`${prefix}layers.0.downsample.`, map, dims[0], dims[1], config.mergeStrides[0]);
+    grid /= config.mergeStrides[0];
+    let tokens = tokensOf(map, dims[1], grid * grid);
+    for (let stage = 1; stage < 4; stage += 1) {
+        const dim = dims[stage];
+        const heads = config.heads[stage];
+        const window = config.windows[stage];
+        for (let i = 0; i < config.depths[stage]; i += 1) {
+            const p = `${prefix}layers.${stage}.blocks.${i}.`;
+            const attended = node('add', [
+                tokens,
+                windowAttention(weights, graph, `${p}attn.`, tokens, dim, heads, window, grid, grid),
+            ]);
+            const local = convBn(`${p}local_conv.`, mapOf(attended, dim, grid, grid), dim, dim, 3, 1, dim);
+            const mixed = tokensOf(local, dim, grid * grid);
+            const inner = gelu(linear(norm(mixed, `${p}mlp.norm.`, dim, EPSILON), `${p}mlp.fc1.`, dim * config.mlpRatio, dim));
+            tokens = node('add', [mixed, linear(inner, `${p}mlp.fc2.`, dim, dim * config.mlpRatio)]);
+        }
+        if (stage < 3) {
+            const stride = config.mergeStrides[stage];
+            const next = dims[stage + 1];
+            const merged = merge(`${prefix}layers.${stage}.downsample.`, mapOf(tokens, dim, grid, grid), dim, next, stride);
+            grid /= stride;
+            tokens = tokensOf(merged, next, grid * grid);
+        }
+    }
+    /* The neck: a pointwise and a 3×3 convolution, each followed by a norm over the channels. */
+    const channelNorm = (x, p, name) => node('reshape', [
+        node('permute', [norm(tokensOf(x, config.neck, grid * grid), p, config.neck, CHANNEL_EPSILON)], { order: [1, 0] }),
+    ], { shape: [config.neck, grid, grid] }, name);
+    const pointwise = node('conv2d', [
+        mapOf(tokens, dims[3], grid, grid),
+        weights.read(`${prefix}neck.0.weight`, [config.neck, dims[3], 1, 1]),
+    ]);
+    const spatial = node('conv2d', [
+        channelNorm(pointwise, `${prefix}neck.1.`),
+        weights.read(`${prefix}neck.2.weight`, [config.neck, config.neck, 3, 3]),
+    ], { padding: 1 });
+    channelNorm(spatial, `${prefix}neck.3.`, output);
+}

package/dist/models/twoWay.d.ts

@@ -0,0 +1,46 @@
+/**
+ * SAM's two-way transformer: a prompt's tokens and an image's, each attending to the other — the
+ * part SAM, MobileSAM and SAM 2 share, under whichever names their checkpoints give it.
+ *
+ * **The tokens attend to themselves, then to the image, then pass an MLP; then the image attends back
+ * to them.** Queries and keys are given their positions afresh at every attention — the tokens as
+ * they were first given, the image by its grid's positions — while values are not. **The first
+ * layer's self-attention replaces the tokens rather than adding to them**, and gives them no
+ * positions, as the upstream's `skip_first_layer_pe` does; every later one is a residual. Attention
+ * between tokens and image narrows to `dim / downsample`. Norms take PyTorch's default ε, 1e-5, and
+ * the MLP's activation is ReLU in every checkpoint here.
+ *
+ * **One definition under two sets of names**, not two definitions: the original's checkpoints say
+ * `out_proj`, `norm1` and `lin1` where Transformers' SAM 2 says `o_proj`, `layer_norm1` and
+ * `proj_in`, and a copy for each would be two implementations of one network, free to drift.
+ */
+import type { GraphBuilder, Weights } from '@driftengine/texture';
+export interface TwoWayConfig {
+    readonly dim: number;
+    readonly heads: number;
+    readonly mlp: number;
+    readonly depth: number;
+    readonly downsample: number;
+}
+/** What the two checkpoint layouts call the transformer's parts. */
+export interface TwoWayNames {
+    /** An attention's output projection. */
+    readonly out: string;
+    /** A layer's `i`-th norm, 1 to 4. */
+    readonly norm: (i: number) => string;
+    readonly finalNorm: string;
+    /** The MLP's two projections. */
+    readonly mlp: readonly [string, string];
+}
+/** The original SAM's names, which MobileSAM keeps. */
+export declare const SAM_NAMES: TwoWayNames;
+/** Transformers' names for SAM 2. */
+export declare const SAM2_NAMES: TwoWayNames;
+/**
+ * The transformer under `prefix` over `tokens`, `[count, dim]`, and `image`, `[cells, dim]` in token
+ * layout with `positions` of the same shape; answers both, attended.
+ */
+export declare function twoWay(weights: Weights, graph: GraphBuilder, prefix: string, config: TwoWayConfig, names: TwoWayNames, tokens: string, image: string, positions: string): {
+    readonly queries: string;
+    readonly keys: string;
+};