export_sam2.py

# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = ["torch==2.9.0", "torchvision==0.24.0", "numpy", "pillow", "huggingface_hub", "hydra-core==1.3.2",
#                 "sam-2 @ git+https://github.com/facebookresearch/sam2.git@2b90b9f5ceec907a1c18123530e92e794ad901a4",
#                 "coreai-torch==0.4.3", "coreai-core==1.0.0b3"]
# ///
"""SAM 2.1 tiny (Meta, Apache-2.0) → cutout-sam2-tiny-<dtype>.aimodel, for the click.

    SAM2_BUILD_CUDA=0 uv run Tools/export_sam2.py [--dtype float16|float32] [--out Tools/exports] [--check image.jpg …]

The network is the authors' own: facebookresearch/sam2 at a pinned commit, with the official SAM 2.1 tiny checkpoint
(facebook/sam2.1-hiera-tiny at a pinned revision). Nothing in it is rewritten; the two entry points are the image
predictor's own steps (SAM2ImagePredictor.set_image and ._predict), traced as they are.

Two entry points, so one encoded picture serves any number of clicks:

  encode  `image` [1,3,1024,1024] float32 RGB in 0…1 (stretched to the square by the host)
          → `image_embedding` [1,256,64,64], `feats_s0` [1,32,256,256], `feats_s1` [1,64,128,128]
          The ImageNet normalisation and SAM 2's no-memory embedding are inside.
  decode  those three, `point` [1,1,2] float32 (pixels of the 1024² square, x then y) and `label` [1,1] float32
          (1 = the object, 0 = not it) → `masks` [1,3,256,256] logits and `scores` [1,3] (predicted IoU).
          The prompt encoder and the mask decoder, multimask output, SAM's padding point added inside.

In a float16 export the weights and arithmetic are half; every input and output stays float32.

--check runs the official SAM2ImagePredictor and the two traced steps on each image (a click at the centre) in
float32 on the CPU and prints the difference, before anything is exported."""
import argparse, sys, time, shutil
from pathlib import Path
import numpy as np, torch, torch.nn as nn

HERE = Path(__file__).resolve().parent
REPOSITORY = "facebook/sam2.1-hiera-tiny"
REVISION = "de431c4043854a71d8101e17995dfe596bf101a5"
CHECKPOINT = "sam2.1_hiera_tiny.pt"
CODE = "facebookresearch/sam2@2b90b9f"
SIZE = 1024

ap = argparse.ArgumentParser()
ap.add_argument("--dtype", default="float16", choices=["float16", "float32"])
ap.add_argument("--out", default=str(HERE / "exports"))
ap.add_argument("--weights", default=None, help="a folder to keep the checkpoint in (default: the HF cache)")
ap.add_argument("--check", nargs="*", default=[], help="images to compare the predictor and the traced steps on")
ap.add_argument("--no-export", action="store_true")
args = ap.parse_args()

from huggingface_hub import hf_hub_download
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

torch.set_num_threads(4)
ckpt = hf_hub_download(REPOSITORY, CHECKPOINT, revision=REVISION, local_dir=args.weights)
sam = build_sam2("configs/sam2.1/sam2.1_hiera_t.yaml", ckpt, device="cpu", mode="eval")
print(f"SAM 2.1 tiny @ {REVISION[:7]}: {sum(p.numel() for p in sam.parameters())/1e6:.1f} M parameters", flush=True)
FEAT_SIZES = [(256, 256), (128, 128), (64, 64)]       # SAM2ImagePredictor._bb_feat_sizes at 1024²
MEAN = torch.tensor([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1)
STD = torch.tensor([0.229, 0.224, 0.225]).reshape(1, 3, 1, 1)


class Encode(nn.Module):
    """SAM2ImagePredictor.set_image after its resize: normalise, backbone, no-memory embedding, reshape."""
    def __init__(self, m, dtype):
        super().__init__(); self.m = m; self.dtype = dtype
        self.register_buffer("mean", MEAN.to(dtype)); self.register_buffer("std", STD.to(dtype))
    def forward(self, image):
        x = (image.to(self.dtype) - self.mean) / self.std
        backbone_out = self.m.forward_image(x)
        _, vision_feats, _, _ = self.m._prepare_backbone_features(backbone_out)
        if self.m.directly_add_no_mem_embed:
            vision_feats[-1] = vision_feats[-1] + self.m.no_mem_embed
        feats = [f.permute(1, 2, 0).reshape(1, -1, *s) for f, s in zip(vision_feats[::-1], FEAT_SIZES[::-1])][::-1]
        return feats[2].to(torch.float32), feats[0].to(torch.float32), feats[1].to(torch.float32)


class Decode(nn.Module):
    """SAM2ImagePredictor._predict for one point, multimask: prompt encoder, mask decoder."""
    def __init__(self, m, dtype):
        super().__init__(); self.m = m; self.dtype = dtype
    def forward(self, image_embedding, feats_s0, feats_s1, point, label):
        d = self.dtype
        # The prompt encoder stays float32 (it is tiny, and its padding point is made float32 inside).
        sparse, dense = self.m.sam_prompt_encoder(points=(point, label.to(torch.int32)), boxes=None, masks=None)
        masks, scores, _, _ = self.m.sam_mask_decoder(
            image_embeddings=image_embedding.to(d), image_pe=self.m.sam_prompt_encoder.get_dense_pe().to(d),
            sparse_prompt_embeddings=sparse.to(d), dense_prompt_embeddings=dense.to(d), multimask_output=True,
            repeat_image=False, high_res_features=[feats_s0.to(d), feats_s1.to(d)])
        return masks.to(torch.float32), scores.to(torch.float32)


def square(path):
    from PIL import Image
    im = Image.open(path).convert("RGB")
    return im, torch.from_numpy(np.asarray(im.resize((SIZE, SIZE), Image.BILINEAR), np.float32) / 255).permute(2, 0, 1)[None]


if args.check:
    predictor = SAM2ImagePredictor(sam)
    enc, dec = Encode(sam, torch.float32).eval(), Decode(sam, torch.float32).eval()
    with torch.no_grad():
        for path in args.check:
            im, x = square(path)
            predictor.set_image(np.asarray(im))
            w, h = im.size
            _, s_ref, low_ref = predictor.predict(point_coords=np.array([[w / 2, h / 2]]), point_labels=np.array([1]),
                                                  multimask_output=True, return_logits=True)
            # The predictor resizes with its own transform; feed the traced steps its tensor so only the steps differ.
            x = predictor._transforms(np.asarray(im))[None]
            x = x * STD + MEAN
            e, s0, s1 = enc(x)
            m, s = dec(e, s0, s1, torch.tensor([[[SIZE / 2, SIZE / 2]]]), torch.ones(1, 1))
            # The predictor clamps its low-resolution logits to ±32; the export does not (the host thresholds).
            dm = np.abs(m[0].clamp(-32, 32).numpy() - low_ref).max()
            print(f"{Path(path).name}: |predictor − traced| masks {dm:.2e} scores {np.abs(s[0].numpy() - s_ref).max():.2e}",
                  flush=True)

if args.no_export:
    sys.exit(0)

import coreai_torch
from coreai.runtime import AIModelAssetMetadata
dtype = torch.float16 if args.dtype == "float16" else torch.float32
sam = sam.to(dtype)
sam.sam_prompt_encoder.float()
enc, dec = Encode(sam, dtype).eval(), Decode(sam, dtype).eval()
t0 = time.time()
converter = coreai_torch.TorchConverter(mode=coreai_torch.TorchConverter.Mode.RELEASE)
with torch.no_grad():
    x = torch.rand(1, 3, SIZE, SIZE)
    ex = torch.export.export(enc, (x,)).run_decompositions(coreai_torch.get_decomp_table())
    converter.add_exported_program(ex, input_names=["image"],
                                   output_names=["image_embedding", "feats_s0", "feats_s1"], entrypoint_name="encode")
    e, s0, s1 = enc(x)
    ex = torch.export.export(dec, (e, s0, s1, torch.tensor([[[300.0, 400.0]]]), torch.ones(1, 1))
                             ).run_decompositions(coreai_torch.get_decomp_table())
    converter.add_exported_program(ex, input_names=["image_embedding", "feats_s0", "feats_s1", "point", "label"],
                                   output_names=["masks", "scores"], entrypoint_name="decode")
    print(f"traced in {time.time()-t0:.0f} s", flush=True)
program = converter.to_coreai()
print(f"converted in {time.time()-t0:.0f} s", flush=True)
meta = AIModelAssetMetadata()
meta.author = "Ravi, Gabeur, Hu, Hu, Ryali, Ma, Khedr, Rädle, Rolland, Gustafson, et al. (Meta) — SAM 2.1 tiny; Core AI export by ObjectCutout"
meta.license = "Apache-2.0"
meta.model_description = (f"SAM 2.1 hiera tiny ({REPOSITORY}@{REVISION[:7]}, code {CODE}), {args.dtype}. "
                          f"'encode': image [1,3,{SIZE},{SIZE}] float32 RGB 0…1 → image_embedding, feats_s0, feats_s1. "
                          f"'decode': those + point [1,1,2] (pixels of the {SIZE}² square) + label [1,1] "
                          f"→ masks [1,3,256,256] logits, scores [1,3].")
meta.creation_date = int(time.time())
out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
asset = out / f"cutout-sam2-tiny-{args.dtype}.aimodel"
if asset.exists(): shutil.rmtree(asset)
program.save_asset(asset, meta)
print(f"saved {asset} ({sum(f.stat().st_size for f in asset.rglob('*'))/1e6:.0f} MB) in {time.time()-t0:.0f} s", flush=True)