File size: 8,799 Bytes
a0b4f4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | # /// 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)
|