| |
| """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 |
|
|
| 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") |
|
|
| |
| |
| |
| 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)], |
| }, |
| } |
| |
| |
| 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"] |
|
|
|
|
| |
| |
| _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 |
|
|
| 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) |
| is_pad = torch.relu(1.0 - d * 1e6) |
| 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) |
|
|
| 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)], |
| ) |
| |
| 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() |
| |
| 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() |
| |
| 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) |
| |
| 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() |
|
|
| |
| gen = torch.Generator().manual_seed(0) |
| pixels = torch.randn(1, 3, args.size, args.size, generator=gen) |
| with torch.no_grad(): |
| |
| 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() |
|
|
