export/sam2-hiera-tiny.py

#!/usr/bin/env python3
"""SAM 2.1 Hiera Tiny -> two ExecuTorch WebGPU programs + prompt tables, 512x512 input.

Same split and I/O contract as the playground's previous export (sam2_encoder.pte, sam2_decoder.pte,
prompt_*.bin); the prompt encoder runs in JS (src/runtime/sam.js) from the tables.

  sam2_encoder.pte  forward(pixels[1,3,512,512] f32: RGB, square resize, /255, ImageNet mean/std)
                    -> (image_embeddings[1,256,32,32], hr0[1,256,128,128], hr1[1,256,64,64])
                    hr0/hr1 are the FPN levels *before* conv_s0/conv_s1 (the decoder applies them).
  sam2_decoder.pte  forward(image_embeddings[1,256,32,32], image_pe[1,256,32,32], sparse[1,1,3,256],
                            dense[1,256,32,32], hr0[1,256,128,128], hr1[1,256,64,64])
                    -> (masks[1,3,128,128] logits, iou_scores[1,1,3], object_score_logits[1,1,1])
                    multimask output (mask tokens 1..3), three prompt tokens (--prompt-tokens): points,
                    box corners, one not-a-point pad, then pads; pads after the first are masked, so
                    clicks, click+exclude and boxes decode exactly as upstream.
  prompt_gaussian.bin          [2,128] f32  random-Fourier matrix (shared_image_embedding)
  prompt_point_embed_bg.bin    [256]   f32  point_embed[0] (negative click)
  prompt_point_embed_fg.bin    [256]   f32  point_embed[1] (positive click)
  prompt_box_tl.bin / _br.bin  [256]   f32  point_embed[2] / [3] (box corners)
  prompt_not_a_point.bin       [256]   f32  not_a_point_embed (padding token)
  prompt_img_pos.bin   [1,256,32,32]   f32  get_image_wide_positional_embeddings() at 512
  prompt_dense.bin     [1,256,32,32]   f32  no_mask_embed broadcast (no mask prompt)

Upstream: the HF transformers port (Apache-2.0) of facebook/sam2.1-hiera-tiny, which is what the
previous export was built from (its decoder weights and all prompt tables are bit-identical to it).
`--weights sam2` uses facebook/sam2-hiera-tiny (SAM 2.0) instead.

Graph changes, all output-preserving: attention goes through the public fused op
torch.ops.et_vk.sdpa (registered as a transformers attention backend); the Hiera windowed absolute
position embedding (bicubic resize of pos_embed + tiled window embedding) for the fixed input size is
computed once in eager and stored as a constant; Hiera's q-pooling of windowed tensors runs as one
batch-1 max_pool2d over the windows stacked along H (exact; keeps every pool in the delegate). Unlike the previous export, the encoder adds
`no_memory_embedding` to the 32x32 embedding, as upstream does for single images
(transformers Sam2Model.forward; `directly_add_no_mem_embed: true` in the official yaml).
`--no-memory-embedding omit` reproduces the previous behaviour.
"""

import argparse
import json
import sys
import types
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, str(Path(__file__).resolve().parent))
import _common  # noqa: E402

WEIGHTS = {
    "sam2.1": {
        "repo": "facebook/sam2.1-hiera-tiny",
        "revision": "de431c4043854a71d8101e17995dfe596bf101a5",
        "files": {
            "model.safetensors": "48c14467e5cf9e51870511feb72c89688e82dd74523142c0538b663e193ac2a7",
        },
    },
    "sam2": {
        "repo": "facebook/sam2-hiera-tiny",
        "revision": "7c218beaf0bb87874785f32b582f640134fc1c09",
        "files": {
            "model.safetensors": "1449035ba6ba2ed524c646da989e5cb3a9b93b97f1c1b2d474e598c3b748d99f",
        },
    },
}
CONFIG_FILES = ("config.json", "preprocessor_config.json", "processor_config.json")

# Validation images. Prompts are in original-image pixels, without padding; site_tokens() turns
# them into the decoder's tokens the way planPrompt() in sam.js does (points, box corners, one
# not-a-point pad, then pads up to the decoder's token count).
IMAGE = {
    "name": "sam-car.png",
    "repo": "huggingface/documentation-images",
    "repo_type": "dataset",
    "revision": "79244b5169b24b4e1444393d9b87349c9182cfa4",
    "file": "transformers/model_doc/sam-car.png",
    "sha256": "c6ae6440e3862dd9768cd9f7098b783ac980629dd7bc1fe69c10b39fca5a84cd",
    "prompts": {
        "door_click": [("pos", 1680, 840)],
        "wheel_click": [("pos", 832, 1088)],
        "door_plus_exclude": [("pos", 1680, 840), ("neg", 1000, 200)],
        "car_box": [("tl", 120, 400), ("br", 2528, 1248)],
        "door_and_wheel_clicks": [("pos", 1680, 840), ("pos", 832, 1088)],
    },
}
# Optional: extra sample images for validation (`--assets-dir`).
# Not public inputs; pinned by sha256 so the numbers are reproducible.
ASSET_IMAGES = {
    "real_pets.jpg": {
        "sha256": "85e4f2f27368a203c30eac68a71e4dde40edac6417b4a8c9df80af5d30b191b2",
        "prompts": {
            "pug_face_click": [("pos", 505, 150)],
            "blanket_click": [("pos", 500, 500)],
            "blanket_minus_face": [("pos", 500, 500), ("neg", 505, 150)],
            "pug_box": [("tl", 165, 65), ("br", 850, 720)],
        },
    },
    "yolo_people.jpg": {
        "sha256": "16d73869e3267a7d4ed00de8e860833bd1657c1b252e94c0c348277adc7b6edb",
        "prompts": {
            "left_man_head_click": [("pos", 570, 260)],
            "right_man_face_click": [("pos", 980, 180)],
            "left_jacket_minus_right": [("pos", 300, 520), ("neg", 880, 520)],
            "right_man_box": [("tl", 750, 40), ("br", 1150, 720)],
        },
    },
}
SIZE = 512
MEAN = (0.485, 0.456, 0.406)
STD = (0.229, 0.224, 0.225)
INPUT_NAMES = ["image_embeddings", "image_pe", "sparse", "dense", "hr0", "hr1"]


# Additive key bias [1, 1, 1, tokens] set by Decoder.forward while the mask decoder runs: it hides
# extra padding tokens (see Decoder.pad_bias). None everywhere else (encoder, references).
_TOKEN_KEY_BIAS = None


def _key_bias(query, key, attention_mask):
    bias = _TOKEN_KEY_BIAS
    if bias is None or key.shape[-2] != bias.shape[-1]:
        return attention_mask
    full = bias.expand(query.shape[0], query.shape[1], query.shape[2], key.shape[-2])
    return full if attention_mask is None else attention_mask + full


def et_vk_attention(
    module, query, key, value, attention_mask, scaling=None, dropout=0.0, **kwargs
):
    """transformers attention backend emitting the fused et_vk.sdpa op (public ExecuTorch Vulkan custom op,
    registered by the WebGPU runtime). q/k/v are [B, H, S, D]; returns [B, S, H, D] like the built-in backends."""
    import executorch.backends.vulkan.custom_ops_lib  # noqa: F401  (defines torch.ops.et_vk.*)

    mask = _key_bias(query, key, attention_mask)
    out = torch.ops.et_vk.sdpa(query, key, value, mask, scaling)
    return out.transpose(1, 2).contiguous(), None


def torch_attention(
    module, query, key, value, attention_mask, scaling=None, dropout=0.0, **kwargs
):
    """Same as et_vk_attention with torch SDPA (for the CPU-only portable decoder)."""
    mask = _key_bias(query, key, attention_mask)
    out = F.scaled_dot_product_attention(query, key, value, attn_mask=mask, scale=scaling)
    return out.transpose(1, 2).contiguous(), None


def register_attention():
    from transformers import AttentionInterface

    AttentionInterface.register("et_vk_sdpa", et_vk_attention)
    AttentionInterface.register("torch_sdpa_padmask", torch_attention)


def site_tokens(prompt, n_tokens):
    """planPrompt() in sam.js for a prompt that fits one decode: points (positive first), box corners,
    one not-a-point pad (SAM 2's own padding), then pads up to the decoder's token count."""
    pos = [t for t in prompt if t[0] == "pos"]
    neg = [t for t in prompt if t[0] == "neg"]
    corners = [t for t in prompt if t[0] in ("tl", "br")]
    native = pos + neg + corners + [("pad",)]
    if len(native) > n_tokens:
        sys.exit(f"prompt {prompt} needs {len(native)} tokens; the decoder takes {n_tokens}")
    return native + [("pad",)] * (n_tokens - len(native))


def snapshot(weights):
    """Pinned upstream snapshot: weights (sha256-checked) + configs at the same revision."""
    w = WEIGHTS[weights]
    path, digest = _common.hf_file(
        w["repo"], "model.safetensors", w["revision"], w["files"]["model.safetensors"]
    )
    for f in CONFIG_FILES:
        _common.hf_file(w["repo"], f, w["revision"])
    return path.parent, digest


def load_model(snap, size, attn="sdpa"):
    from transformers import Sam2Config, Sam2Model

    cfg = Sam2Config.from_pretrained(snap)
    if size != cfg.image_size:
        f = size // cfg.prompt_encoder_config.patch_size
        cfg.image_size = size
        cfg.prompt_encoder_config.image_size = size
        cfg.vision_config.backbone_config.image_size = [size, size]
        cfg.vision_config.backbone_feature_sizes = [
            [4 * f, 4 * f],
            [2 * f, 2 * f],
            [f, f],
        ]
    model, info = Sam2Model.from_pretrained(
        snap, config=cfg, attn_implementation=attn, output_loading_info=True
    )
    if info["missing_keys"] or info["mismatched_keys"] or info["unexpected_keys"]:
        sys.exit(f"weights did not load cleanly: {info}")
    return model.eval()


def pool_windows_as_one_image():
    """Hiera's q-pooling (2x2 max pool, stride 2) on windowed tensors [nW, h, w, C] is rewritten to pool
    the windows stacked along H as a single [1, C, nW*h, w] image. h is a multiple of the stride, so no
    pooling cell straddles two windows and the result is identical; batch-1 pooling is what the Vulkan
    partitioner can represent (a [256, 384, 4, 4] pool is left on CPU otherwise)."""
    from transformers.models.sam2 import modeling_sam2

    original = getattr(modeling_sam2, "_upstream_do_pool", modeling_sam2.do_pool)
    modeling_sam2._upstream_do_pool = original

    def do_pool(x, query_stride=None):
        if query_stride is None or x.shape[0] == 1:
            return original(x, query_stride)
        b, h, w, c = x.shape
        sh, sw = (query_stride, query_stride) if isinstance(query_stride, int) else tuple(query_stride)
        if h % sh or w % sw:
            return original(x, query_stride)
        y = x.reshape(1, b * h, w, c).permute(0, 3, 1, 2)
        y = F.max_pool2d(y, kernel_size=query_stride, stride=query_stride, ceil_mode=False)
        return y.permute(0, 2, 3, 1).reshape(b, h // sh, w // sw, c)

    modeling_sam2.do_pool = do_pool


def upstream_pool():
    """Undo pool_windows_as_one_image() (upstream code path, for references)."""
    from transformers.models.sam2 import modeling_sam2

    modeling_sam2.do_pool = getattr(modeling_sam2, "_upstream_do_pool", modeling_sam2.do_pool)


def freeze_hiera_pos_embed(model, size):
    bb = model.vision_encoder.backbone
    with torch.no_grad():
        frozen = bb._get_pos_embed((size // 4, size // 4)).detach().contiguous().clone()
    bb.register_buffer("frozen_pos_embed", frozen, persistent=False)
    bb._get_pos_embed = types.MethodType(lambda self, _hw: self.frozen_pos_embed, bb)
    return model


class Encoder(torch.nn.Module):
    def __init__(self, model, add_no_memory=True):
        super().__init__()
        self.vision = model.vision_encoder
        self.add_no_memory = add_no_memory
        self.register_buffer(
            "no_mem",
            model.no_memory_embedding.detach().reshape(1, -1, 1, 1).clone(),
            persistent=False,
        )

    def forward(self, pixels):
        hr0, hr1, emb = self.vision(pixels).fpn_hidden_states
        if self.add_no_memory:
            emb = emb + self.no_mem
        return emb, hr0, hr1


class Decoder(torch.nn.Module):
    """Mask decoder on precomputed prompt embeddings.

    A pad (not-a-point) token that directly follows another pad is hidden from every attention as a
    key (additive -1e9, which underflows to an exact 0 weight). Nothing reads such a token's own
    output, so a fixed-size decode of [pos, pad, pad] equals upstream's [pos, pad], and [tl, br, pad]
    equals upstream's box prompt. The first pad is kept: it is SAM 2's own padding token."""

    def __init__(self, model):
        super().__init__()
        self.md = model.mask_decoder
        self.n_out = self.md.obj_score_token.weight.shape[0] + 1 + self.md.num_mask_tokens
        self.register_buffer(
            "not_a_point",
            model.prompt_encoder.not_a_point_embed.weight[0].detach().clone(),
            persistent=False,
        )

    def pad_bias(self, sparse):
        d = (sparse - self.not_a_point).abs().sum(-1)  # [1, 1, N]; 0 only for an exact pad token
        is_pad = torch.relu(1.0 - d * 1e6)  # 1.0 for a pad, 0.0 otherwise
        repeated = is_pad[..., 1:] * is_pad[..., :-1]
        keep = torch.zeros(1, 1, self.n_out + 1, dtype=sparse.dtype)
        return (torch.cat([keep, repeated], dim=-1) * -1e9).unsqueeze(2)  # [1, 1, 1, n_out + N]

    def forward(self, emb, image_pe, sparse, dense, hr0, hr1):
        global _TOKEN_KEY_BIAS
        _TOKEN_KEY_BIAS = self.pad_bias(sparse) if sparse.shape[2] > 1 else None
        try:
            return self._decode(emb, image_pe, sparse, dense, hr0, hr1)
        finally:
            _TOKEN_KEY_BIAS = None

    def _decode(self, emb, image_pe, sparse, dense, hr0, hr1):
        masks, iou, _, obj = self.md(
            image_embeddings=emb,
            image_positional_embeddings=image_pe,
            sparse_prompt_embeddings=sparse,
            dense_prompt_embeddings=dense,
            multimask_output=True,
            high_resolution_features=[self.md.conv_s0(hr0), self.md.conv_s1(hr1)],
        )
        # [B, point_batch=1, 3, H, W] -> [B, 3, H, W], as the previous export.
        return masks[:, 0], iou, obj


def prompt_tables(model):
    pe = model.prompt_encoder
    with torch.no_grad():
        dense = (
            pe.no_mask_embed.weight.reshape(1, -1, 1, 1)
            .expand(1, -1, *pe.image_embedding_size)
            .contiguous()
        )
        tabs = {
            "gaussian": model.shared_image_embedding.positional_embedding,
            "point_embed_bg": pe.point_embed.weight[0],
            "point_embed_fg": pe.point_embed.weight[1],
            "box_tl": pe.point_embed.weight[2],
            "box_br": pe.point_embed.weight[3],
            "not_a_point": pe.not_a_point_embed.weight[0],
            "img_pos": model.get_image_wide_positional_embeddings(),
            "dense": dense,
        }
    return {
        k: v.detach().to(torch.float32).contiguous().numpy() for k, v in tabs.items()
    }


def encode_tokens(tokens, width, height, tabs, size=SIZE):
    """Python port of the playground's browser-side prompt encoding (float64 math, f32 out)."""
    g = tabs["gaussian"].reshape(-1).astype(np.float64)
    hidden = tabs["not_a_point"].size
    D = hidden // 2
    kind = {
        "neg": tabs["point_embed_bg"],
        "pos": tabs["point_embed_fg"],
        "tl": tabs["box_tl"],
        "br": tabs["box_br"],
    }
    out = np.zeros((len(tokens), hidden), np.float64)
    for r, t in enumerate(tokens):
        if t[0] == "pad":
            out[r] = tabs["not_a_point"]
            continue
        x, y = round(t[1] * 10) / 10, round(t[2] * 10) / 10
        cx = (2 * (x * (size / width) + 0.5)) / size - 1
        cy = (2 * (y * (size / height) + 0.5)) / size - 1
        a = 2 * np.pi * (cx * g[:D] + cy * g[D:])
        out[r, :D] = np.sin(a) + kind[t[0]][:D]
        out[r, D:] = np.cos(a) + kind[t[0]][D:]
    return torch.from_numpy(out.astype(np.float32)).reshape(1, 1, len(tokens), hidden)


def load_image(path, size):
    from PIL import Image

    img = Image.open(path).convert("RGB")
    a = np.asarray(img.resize((size, size), Image.BILINEAR), dtype=np.float32) / 255.0
    a = (a - np.array(MEAN, np.float32)) / np.array(STD, np.float32)
    return torch.from_numpy(a.transpose(2, 0, 1).copy()).unsqueeze(0), img


def iou(a, b):
    a, b = np.asarray(a) > 0, np.asarray(b) > 0
    u = np.logical_or(a, b).sum()
    return round(float(np.logical_and(a, b).sum() / u), 4) if u else 1.0


def upsample(masks, w, h):
    """Bilinear, align_corners=False: the same mapping as upsampleMask() in sam.js."""
    t = torch.from_numpy(np.ascontiguousarray(masks, np.float32)).reshape(
        1, -1, *np.asarray(masks).shape[-2:]
    )
    return F.interpolate(t, size=(h, w), mode="bilinear", align_corners=False)[
        0
    ].numpy()


def upstream_inputs(tokens):
    """The same prompt in transformers' native form (Sam2PromptEncoder pads points with a not-a-point
    token and always adds one after box corners, like the original SAM 2 predictor)."""
    pts = [t for t in tokens if t[0] in ("pos", "neg")]
    corners = [t for t in tokens if t[0] in ("tl", "br")]
    kw = {}
    if corners:
        (_, x0, y0), (_, x1, y1) = corners
        kw["input_boxes"] = [[[x0, y0, x1, y1]]]
    if pts:
        kw["input_points"] = [[[[t[1], t[2]] for t in pts]]]
        kw["input_labels"] = [[[1 if t[0] == "pos" else 0 for t in pts]]]
    return kw


def upstream_reference(model, processor, img, tokens):
    """Upstream transformers pipeline end to end (Sam2Processor -> Sam2Model -> post_process_masks):
    mask logits at the original image size, and the IoU head's scores."""
    inputs = processor(images=img, return_tensors="pt", **upstream_inputs(tokens))
    with torch.no_grad():
        o = model(
            pixel_values=inputs["pixel_values"],
            input_points=inputs.get("input_points"),
            input_labels=inputs.get("input_labels"),
            input_boxes=inputs.get("input_boxes"),
            multimask_output=True,
        )
    m = processor.post_process_masks(
        o.pred_masks, inputs["original_sizes"], binarize=False
    )[0]
    return m[0].numpy(), o.iou_scores[0, 0].numpy()


def upstream_512_same_pixels(model, pixels, tokens, w, h):
    """Upstream Sam2Model.forward (512 config) on exactly the pixels the export sees."""
    s = model.prompt_encoder.input_image_size
    kw = upstream_inputs(tokens)
    sx, sy = s / w, s / h
    args = {}
    if "input_boxes" in kw:
        x0, y0, x1, y1 = kw["input_boxes"][0][0]
        args["input_boxes"] = torch.tensor(
            [[[x0 * sx, y0 * sy, x1 * sx, y1 * sy]]], dtype=torch.float32
        )
    if "input_points" in kw:
        args["input_points"] = torch.tensor(
            [[[[x * sx, y * sy] for x, y in kw["input_points"][0][0]]]],
            dtype=torch.float32,
        )
        args["input_labels"] = torch.tensor(kw["input_labels"], dtype=torch.int64)
    with torch.no_grad():
        o = model(pixel_values=pixels, multimask_output=True, **args)
    return o.pred_masks[0, 0].numpy(), o.iou_scores[0, 0].numpy()


CONTRACT = {
    "sam2_encoder": (
        [[1, 3, SIZE, SIZE]],
        [[1, 256, SIZE // 16, SIZE // 16], [1, 256, SIZE // 4, SIZE // 4], [1, 256, SIZE // 8, SIZE // 8]],
    ),
    "sam2_decoder": (
        [[1, 256, SIZE // 16, SIZE // 16], [1, 256, SIZE // 16, SIZE // 16], [1, 1, "N", 256],
         [1, 256, SIZE // 16, SIZE // 16], [1, 256, SIZE // 4, SIZE // 4], [1, 256, SIZE // 8, SIZE // 8]],
        [[1, 3, SIZE // 4, SIZE // 4], [1, 1, 3], [1, 1, 1]],
    ),
}


def program_io(pte_bytes):
    """Method names and forward() I/O (shape, dtype, dim order) as serialized in the program."""
    from executorch.exir._serialize._program import deserialize_pte_binary

    pte = deserialize_pte_binary(pte_bytes)
    prog = pte.program if hasattr(pte, "program") else pte

    def desc(plan, i):
        v = plan.values[i].val
        return {"shape": list(v.sizes), "dtype": v.scalar_type.name, "dim_order": list(v.dim_order)}

    plan = {p.name: p for p in prog.execution_plan}
    f = plan["forward"]
    return {
        "methods": sorted(plan),
        "inputs": [desc(f, i) for i in f.inputs],
        "outputs": [desc(f, i) for i in f.outputs],
    }


def check_contract(name, io, size, n_tokens):
    """The playground's programs: one `forward`, fp32, contiguous NCHW, these shapes."""
    if size != SIZE:
        return "skipped (non-default size)"
    want_in, want_out = CONTRACT[name]
    want_in = [[n_tokens if d == "N" else d for d in s] for s in want_in]
    ok = (
        io["methods"] == ["forward"]
        and [t["shape"] for t in io["inputs"]] == want_in
        and [t["shape"] for t in io["outputs"]] == want_out
        and all(t["dtype"] == "FLOAT" and t["dim_order"] == list(range(len(t["shape"]))) for t in io["inputs"] + io["outputs"])
    )
    if not ok:
        sys.exit(f"{name}: I/O does not match the previous contract: {json.dumps(io)}")
    return "matches the previous export (shapes, fp32, contiguous dim order, single forward method)"


def write_if_changed(path, data):
    """Leave an identical file alone; otherwise replace it with a new inode, so hard links made to an
    earlier output elsewhere keep their content."""
    data = bytes(data)
    if path.exists():
        if path.read_bytes() == data:
            return
        path.unlink()
    path.write_bytes(data)


def export_programs(enc, dec, enc_std, dec_std, pixels, dec_inputs):
    with torch.no_grad():
        ep_enc = torch.export.export(enc, (pixels,), strict=False)
        ep_dec = torch.export.export(dec, dec_inputs, strict=False)
        ep_enc_std = torch.export.export(enc_std, (pixels,), strict=False)
        ep_dec_std = torch.export.export(dec_std, dec_inputs, strict=False)
        return {
            "sam2_encoder": (
                ep_enc,
                _common.lower_webgpu(ep_enc),
                _common.lower_portable(ep_enc_std),
            ),
            "sam2_decoder": (
                ep_dec,
                _common.lower_webgpu(ep_dec),
                _common.lower_portable(ep_dec_std),
            ),
        }


def validation_images(args):
    from huggingface_hub import hf_hub_download

    out = []
    p = Path(
        hf_hub_download(
            IMAGE["repo"],
            IMAGE["file"],
            repo_type=IMAGE["repo_type"],
            revision=IMAGE["revision"],
        )
    )
    digest = _common.sha256_file(p)
    if digest != IMAGE["sha256"]:
        sys.exit(f"{IMAGE['file']}: sha256 {digest} != {IMAGE['sha256']}")
    out.append(
        (
            IMAGE["name"],
            p,
            digest,
            IMAGE["prompts"],
            f"hf://datasets/{IMAGE['repo']}@{IMAGE['revision']}/{IMAGE['file']}",
        )
    )
    if args.assets_dir:
        for name, spec in ASSET_IMAGES.items():
            ap = Path(args.assets_dir) / name
            d = _common.sha256_file(ap)
            if spec["sha256"] and d != spec["sha256"]:
                sys.exit(f"{ap}: sha256 {d} != {spec['sha256']}")
            out.append(
                (
                    name,
                    ap,
                    d,
                    spec["prompts"],
                    f"sample image {name}",
                )
            )
    return out


def validate(args, snap, model, tabs, progs, enc, dec, enc_std):
    from transformers import Sam2Processor

    val = {
        "cpu_runtime": "executorch pybindings (portable_lib) from the executorch 1.5.1 wheel"
    }
    ep_enc, wg_enc, port_enc = progs["sam2_encoder"]
    ep_dec, wg_dec, port_dec = progs["sam2_decoder"]
    img_pe = torch.from_numpy(tabs["img_pos"])
    dense = torch.from_numpy(tabs["dense"])

    ref512 = load_model(snap, args.size)
    ref1024 = load_model(snap, 1024)
    processor = Sam2Processor.from_pretrained(snap)
    other = Encoder(model, add_no_memory=not enc.add_no_memory).eval()
    # Everything below runs the upstream do_pool; the export used the rewrite (checked to be exact).
    upstream_pool()

    for p_name, ep, wg in (("encoder", ep_enc, wg_enc), ("decoder", ep_dec, wg_dec)):
        _, err = _common.try_run_webgpu_pte_on_cpu(wg.buffer, ep.example_inputs[0])
        val[f"webgpu_{p_name}_pte_on_cpu"] = (
            f"not runnable on CPU pybindings: {err}" if err else "ran"
        )

    val["images"] = {}
    refs = Path(args.refs_dir) if args.refs_dir else None
    refs_meta = {"prompt_tokens": args.prompt_tokens, "prompts": {}}
    for name, path, digest, prompts, src in validation_images(args):
        pixels, img = load_image(path, args.size)
        W, H = img.size
        with torch.no_grad():
            ref = enc(pixels)
            pool_windows_as_one_image()
            rewritten = enc(pixels)
            upstream_pool()
            exported = ep_enc.module()(pixels)
            std = enc_std(pixels)
            emb_other = other(pixels)[0]
        port = _common.run_pte(port_enc.buffer, (pixels,))
        rec = {
            "source": src,
            "sha256": digest,
            "size": [W, H],
            "encoder": {},
            "prompts": {},
        }
        for i, n in enumerate(["image_embeddings", "hr0", "hr1"]):
            rec["encoder"][n] = {
                "exported_vs_eager": _common.compare(
                    exported[i].numpy(), ref[i].numpy()
                ),
                "pool_rewrite_vs_upstream_code_max_abs": float((rewritten[i] - ref[i]).abs().max()),
                "eager_torch_sdpa_vs_eager_et_vk_sdpa": _common.compare(
                    std[i].numpy(), ref[i].numpy()
                ),
                "portable_pte_vs_eager": _common.compare(port[i], ref[i].numpy()),
            }
        emb, hr0, hr1 = (t.contiguous() for t in ref)
        if refs is not None and name == IMAGE["name"]:
            refs.mkdir(parents=True, exist_ok=True)
            refs.joinpath("pixels.f32").write_bytes(pixels.numpy().tobytes())
            for i, t in enumerate((emb, hr0, hr1)):
                refs.joinpath(f"enc_out{i}.f32").write_bytes(t.numpy().tobytes())
            refs_meta.update(image=src, image_sha256=digest, size=[W, H])
        for pname, toks in prompts.items():
            dtoks = site_tokens(toks, args.prompt_tokens)
            sparse = encode_tokens(dtoks, W, H, tabs, args.size)
            ins = (emb, img_pe, sparse, dense, hr0, hr1)
            with torch.no_grad():
                eager = [t.numpy() for t in dec(*ins)]
                exp = [t.numpy() for t in ep_dec.module()(*ins)]
                alt_out = dec(emb_other, img_pe, sparse, dense, hr0, hr1)
                alt, alt_iou = alt_out[0].numpy(), alt_out[1].numpy()
            # End to end through both portable programs (encoder output -> decoder input).
            pd = _common.run_pte(
                port_dec.buffer,
                (
                    torch.from_numpy(np.ascontiguousarray(port[0])),
                    img_pe,
                    sparse,
                    dense,
                    torch.from_numpy(np.ascontiguousarray(port[1])),
                    torch.from_numpy(np.ascontiguousarray(port[2])),
                ),
            )
            m_pte = pd[0][0]
            u512_m, u512_iou = upstream_512_same_pixels(ref512, pixels, toks, W, H)
            u_m, u_iou = upstream_reference(ref1024, processor, img, toks)
            up = upsample(m_pte, W, H)
            up_alt = upsample(alt[0], W, H)
            best = int(np.argmax(u_iou))
            own = int(np.argmax(np.asarray(pd[1]).reshape(-1)))
            alt_own = int(np.argmax(alt_iou.reshape(-1)))
            if refs is not None and name == IMAGE["name"]:
                for i, t in enumerate(eager):
                    refs.joinpath(f"dec_{pname}_out{i}.f32").write_bytes(np.ascontiguousarray(t, np.float32).tobytes())
                refs.joinpath(f"sparse_{pname}.f32").write_bytes(sparse.numpy().tobytes())
                refs.joinpath(f"up512_{pname}_masks.f32").write_bytes(np.ascontiguousarray(u512_m, np.float32).tobytes())
                refs_meta["prompts"][pname] = {"prompt": [list(t) for t in toks], "tokens": [list(t) for t in dtoks]}
            rec["prompts"][pname] = {
                "prompt": [list(t) for t in toks],
                "decoder_tokens": [t[0] for t in dtoks],
                "eager_vs_upstream_512_masks": _common.compare(eager[0][0], u512_m),
                "eager_vs_upstream_512_iou_scores": _common.compare(eager[1], u512_iou),
                "exported_vs_eager_masks": _common.compare(exp[0], eager[0]),
                "portable_e2e_vs_eager_masks": _common.compare(pd[0], eager[0]),
                "portable_e2e_vs_eager_iou_scores": _common.compare(pd[1], eager[1]),
                "portable_e2e_vs_eager_object_score": _common.compare(pd[2], eager[2]),
                "mask_iou_portable_vs_eager": [
                    iou(m_pte[i], eager[0][0][i]) for i in range(3)
                ],
                "mask_iou_portable_vs_upstream_512_same_pixels": [
                    iou(m_pte[i], u512_m[i]) for i in range(3)
                ],
                "mask_iou_portable_vs_upstream_1024_pipeline_at_image_res": [
                    iou(up[i], u_m[i]) for i in range(3)
                ],
                "portable_iou_scores": np.asarray(pd[1]).reshape(-1).round(4).tolist(),
                "upstream_512_iou_scores": u512_iou.round(4).tolist(),
                "upstream_1024_iou_scores": u_iou.round(4).tolist(),
                "upstream_1024_best_mask": best,
                "best_mask_iou_portable_vs_upstream_1024": iou(up[best], u_m[best]),
                "portable_object_score_logit": round(
                    float(np.asarray(pd[2]).reshape(-1)[0]), 4
                ),
                "mask_area_fraction_portable": [
                    round(float((m_pte[i] > 0).mean()), 4) for i in range(3)
                ],
                "upstream_1024_best_mask_vs_closest_export_mask": max(iou(up[i], u_m[best]) for i in range(3)),
                "export_best_by_own_iou_head_vs_upstream_1024_best": iou(up[own], u_m[best]),
                "variant_%s_no_memory_embedding" % ("without" if enc.add_no_memory else "with"): {
                    "mask_iou_vs_this_export": [iou(alt[0][i], eager[0][0][i]) for i in range(3)],
                    "mask_iou_vs_upstream_1024_pipeline_at_image_res": [iou(up_alt[i], u_m[i]) for i in range(3)],
                    "upstream_1024_best_mask_vs_closest_mask": max(iou(up_alt[i], u_m[best]) for i in range(3)),
                    "best_by_own_iou_head_vs_upstream_1024_best": iou(up_alt[alt_own], u_m[best]),
                },
            }
        val["images"][name] = rec
        print(json.dumps({name: rec}, indent=1), flush=True)
    if refs is not None:
        refs.joinpath("meta.json").write_text(json.dumps(refs_meta, indent=1) + "\n")

    if args.internal_dir:
        d = Path(args.internal_dir)
        val["prompt_tables_vs_internal"] = {
            f"prompt_{k}.bin": "bit-identical"
            if (d / f"prompt_{k}.bin").read_bytes() == v.astype(np.float32).tobytes()
            else "max abs %.3g"
            % float(
                np.abs(
                    np.fromfile(d / f"prompt_{k}.bin", dtype=np.float32) - v.reshape(-1)
                ).max()
            )
            for k, v in tabs.items()
        }
    return val


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--out", required=True)
    ap.add_argument("--weights", choices=sorted(WEIGHTS), default="sam2.1")
    ap.add_argument(
        "--size", type=int, default=SIZE, help="square input size (multiple of 32)"
    )
    ap.add_argument(
        "--no-memory-embedding",
        choices=["add", "omit"],
        default="add",
        help="add (upstream single-image behaviour, default) or omit (as the previous export)",
    )
    ap.add_argument(
        "--assets-dir",
        help="optional: directory of sample images for extra validation",
    )
    ap.add_argument(
        "--internal-dir",
        help="optional: compare prompt tables with an existing prompt_*.bin set",
    )
    ap.add_argument(
        "--prompt-tokens",
        type=int,
        default=3,
        help="sparse prompt tokens per decode (trailing extra pads are masked, so any prompt that fits matches upstream)",
    )
    ap.add_argument(
        "--refs-dir",
        help="optional: write CPU reference tensors for the browser harness (sam-car.png prompts)",
    )
    ap.add_argument(
        "--skip-validation",
        action="store_true",
        help="export only (for the byte-identical re-run check)",
    )
    args = ap.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    torch.manual_seed(0)

    snap, digest = snapshot(args.weights)
    register_attention()
    pool_windows_as_one_image()
    add_nm = args.no_memory_embedding == "add"
    model = freeze_hiera_pos_embed(
        load_model(snap, args.size, attn="et_vk_sdpa"), args.size
    )
    enc, dec = Encoder(model, add_nm).eval(), Decoder(model).eval()
    tabs = prompt_tables(model)
    # Same weights with torch SDPA: only for the CPU portable .pte (et_vk.sdpa has no portable kernel).
    model_std = freeze_hiera_pos_embed(
        load_model(snap, args.size, attn="sdpa"), args.size
    )
    model_std_dec = load_model(snap, args.size, attn="torch_sdpa_padmask")
    enc_std, dec_std = Encoder(model_std, add_nm).eval(), Decoder(model_std_dec).eval()

    # Example inputs for export: a fixed-seed tensor (the graph does not depend on values).
    gen = torch.Generator().manual_seed(0)
    pixels = torch.randn(1, 3, args.size, args.size, generator=gen)
    with torch.no_grad():
        # Contiguous example inputs: torch.export records each input's dim order.
        emb, hr0, hr1 = (t.contiguous() for t in enc(pixels))
    sparse = encode_tokens(
        site_tokens([("pos", args.size / 2, args.size / 2)], args.prompt_tokens),
        args.size,
        args.size,
        tabs,
        args.size,
    )
    dec_inputs = (
        emb,
        torch.from_numpy(tabs["img_pos"]),
        sparse,
        torch.from_numpy(tabs["dense"]),
        hr0,
        hr1,
    )
    progs = export_programs(enc, dec, enc_std, dec_std, pixels, dec_inputs)

    files, graphs = {}, {}
    for name, (_, wg, port) in progs.items():
        write_if_changed(out / f"{name}.pte", wg.buffer)
        write_if_changed(out / f"{name}_portable.pte", port.buffer)
        graphs[name] = _common.delegate_summary(wg.buffer)
        graphs[name]["io"] = program_io(wg.buffer)
        graphs[name]["io_contract"] = check_contract(name, graphs[name]["io"], args.size, args.prompt_tokens)
        if graphs[name]["non_delegated_kernels"] or len(graphs[name]["delegates"]) != 1:
            sys.exit(f"{name}: expected one delegate and no CPU kernels: {json.dumps(graphs[name])[:2000]}")
    for k, v in tabs.items():
        write_if_changed(out / f"prompt_{k}.bin", v.astype(np.float32).tobytes())
    for p in sorted(out.iterdir()):
        if p.suffix in (".pte", ".bin"):
            files[p.name] = {
                "bytes": p.stat().st_size,
                "sha256": _common.sha256_file(p),
            }
            if p.name.endswith("_portable.pte"):
                files[p.name]["note"] = (
                    "CPU-only validation artifact (portable kernels, no delegate; torch SDPA instead of et_vk.sdpa); not for release"
                )

    w = WEIGHTS[args.weights]
    q = args.size // 4
    recipe = {
        "model": f"{w['repo'].split('/')[1]} ({args.size} encoder/decoder split)",
        "upstream": {
            "model": w["repo"],
            "revision": w["revision"],
            "official": True,
            "license": "Apache-2.0",
            "files": {"model.safetensors": {"sha256": digest}},
            "note": "transformers Sam2Model port published by the SAM 2 authors on the Hub",
        },
        "versions": _common.versions(),
        "export": {
            "api": "torch.export.export(strict=False) -> to_edge_transform_and_lower(partitioner=[WebGPUPartitioner()]) -> to_executorch(), one program per module",
            "partitioner": "executorch.backends.webgpu.partitioner.WebGPUPartitioner() (default compile options; wraps VulkanPartitioner)",
            "model_code": "transformers Sam2Model; attention through torch.ops.et_vk.sdpa (transformers attention backend, recipes/sam2-hiera-tiny.py:et_vk_attention); config image_size overridden 1024 -> %d"
            % args.size,
            "programs": {
                "sam2_encoder.pte": {
                    "method": "forward",
                    "inputs": [
                        {
                            "name": "pixels",
                            "shape": [1, 3, args.size, args.size],
                            "dtype": "float32",
                        }
                    ],
                    "outputs": [
                        {
                            "name": "image_embeddings",
                            "shape": list(emb.shape),
                            "dtype": "float32",
                            "note": "no_memory_embedding %s"
                            % ("added" if add_nm else "not added"),
                        },
                        {
                            "name": "hr0",
                            "shape": list(hr0.shape),
                            "dtype": "float32",
                            "note": "before mask_decoder.conv_s0",
                        },
                        {
                            "name": "hr1",
                            "shape": list(hr1.shape),
                            "dtype": "float32",
                            "note": "before mask_decoder.conv_s1",
                        },
                    ],
                },
                "sam2_decoder.pte": {
                    "method": "forward",
                    "prompt_tokens": args.prompt_tokens,
                    "inputs": [
                        {"name": n, "shape": list(t.shape), "dtype": "float32"}
                        for n, t in zip(INPUT_NAMES, dec_inputs)
                    ],
                    "outputs": [
                        {
                            "name": "masks",
                            "shape": [1, 3, q, q],
                            "dtype": "float32",
                            "note": "logits, mask tokens 1..3",
                        },
                        {"name": "iou_scores", "shape": [1, 1, 3], "dtype": "float32"},
                        {
                            "name": "object_score_logits",
                            "shape": [1, 1, 1],
                            "dtype": "float32",
                        },
                    ],
                },
            },
            "prompt_tables": {
                f"prompt_{k}.bin": {
                    "shape": list(v.shape),
                    "dtype": "float32 little-endian",
                }
                for k, v in tabs.items()
            },
            "prompt_table_sources": {
                "prompt_gaussian.bin": "shared_image_embedding.positional_embedding",
                "prompt_point_embed_bg.bin": "prompt_encoder.point_embed.weight[0]",
                "prompt_point_embed_fg.bin": "prompt_encoder.point_embed.weight[1]",
                "prompt_box_tl.bin": "prompt_encoder.point_embed.weight[2]",
                "prompt_box_br.bin": "prompt_encoder.point_embed.weight[3]",
                "prompt_not_a_point.bin": "prompt_encoder.not_a_point_embed.weight[0]",
                "prompt_img_pos.bin": "Sam2Model.get_image_wide_positional_embeddings() with image_size=%d"
                % args.size,
                "prompt_dense.bin": "prompt_encoder.no_mask_embed.weight broadcast to [1,256,%d,%d]"
                % (args.size // 16, args.size // 16),
            },
            "graph_rewrites": [
                "Hiera q-pooling of windowed tensors [nW,h,w,C] done as one batch-1 max_pool2d over the windows stacked along H (exact: h is a multiple of the stride); keeps all pools in the delegate",
                "attention via torch.ops.et_vk.sdpa (same math as torch SDPA; checked in validation.json)",
                "decoder: a not-a-point token directly after another one gets an additive -1e9 key bias in every attention (computed in-graph from the sparse input), so padding up to the fixed token count does not change the result",
                "Hiera position embedding (bicubic pos_embed resize + tiled window embedding) for the fixed input size computed in eager and stored as a constant",
            ],
            "differences_from_upstream_pipeline": [
                "input %dx%d instead of 1024x1024 (masks %dx%d instead of 256x256)"
                % (args.size, args.size, q, q),
                "fixed %d prompt tokens per decode: points, box corners, one not-a-point pad, then padding pads that the decoder masks out (exactly upstream's prompt for any prompt that fits; larger prompts are composed by the site)" % args.prompt_tokens,
                "no mask prompt input (dense = no_mask_embed)",
            ],
            "preprocessing_outside_graph": "RGB, resize to %dx%d without keeping aspect, /255, ImageNet mean/std"
            % (args.size, args.size),
            "dtype": "float32",
            "quantization": "none (fp32)",
        },
        "files": files,
        "graph": graphs,
        "command": _common.command_line(out),
    }
    (out / "recipe.json").write_text(json.dumps(recipe, indent=2) + "\n")
    print(json.dumps({"files": files, "graph": graphs}, indent=1), flush=True)
    if args.skip_validation:
        return
    val = validate(args, snap, model, tabs, progs, enc, dec, enc_std)
    (out / "validation.json").write_text(json.dumps(val, indent=2) + "\n")
    print(json.dumps(val, indent=1))


if __name__ == "__main__":
    main()