fp16 build of EfficientViT-SAM

fp16/config.json

@@ -0,0 +1,26 @@
+{
+ "model_type": "efficientvit_sam",
+ "variant": "l2",
+ "image_size": 1024,
+ "encoder_size": 512,
+ "width_list": [
+  32,
+  64,
+  128,
+  256,
+  512
+ ],
+ "depth_list": [
+  1,
+  2,
+  2,
+  8,
+  8
+ ],
+ "qkv_dim": 32,
+ "neck_depth": 12,
+ "neck_expand": 1,
+ "prompt_dim": 256,
+ "mask_in_chans": 16,
+ "precision": "fp16"
+}

fp16/mlx_model.py

@@ -0,0 +1,572 @@
+"""EfficientViT-SAM (mit-han-lab/efficientvit, Apache-2.0) in MLX: a port of `efficientvit_sam_l2` (EfficientViT-L2
+image encoder at 512 px with its SAM neck, then SAM's prompt encoder and mask decoder) and EfficientViTSamPredictor.
+
+    import mlx_model
+    model = mlx_model.load("bf16")                        # a folder with config.json + model.safetensors
+    emb = mlx_model.embed(model, rgb)                     # rgb HxWx3 uint8 -> image embedding (+ sizes)
+    masks, scores, low_res, logits = mlx_model.predict(model, emb, points=[[x, y]], labels=[1], multimask=True)
+
+Runs on its own next to its weights (mlx, numpy, pillow). NHWC; BatchNorms folded into their convs at conversion;
+1x1 convolutions are linear layers (so 8bit / 6bit / 4bit quantize them); LiteMLA's ReLU linear attention and the
+neck's bicubic upsampling run in float32 as in the original.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+CONFIG = {"model_type": "efficientvit_sam", "variant": "l2", "image_size": 1024, "encoder_size": 512,
+          "width_list": [32, 64, 128, 256, 512], "depth_list": [1, 2, 2, 8, 8], "qkv_dim": 32,
+          "neck_depth": 12, "neck_expand": 1, "prompt_dim": 256, "mask_in_chans": 16}
+PRECISIONS = ["bf16", "fp16"]                      # 8bit / 6bit / 4bit don't work: model.toml
+PIXEL_MEAN = np.array([123.675, 116.28, 103.53], np.float32)
+PIXEL_STD = np.array([58.395, 57.12, 57.375], np.float32)
+GROUP = 64
+
+
+
+
+def quantizable(path: str, module) -> bool:
+    return isinstance(module, nn.Linear) and module.weight.shape[-1] % GROUP == 0
+
+
+def resize(x, size):
+    """torch F.interpolate(mode="bilinear", align_corners=False), NHWC."""
+    b, h, w, c = x.shape
+    oh, ow = int(size[0]), int(size[1])
+    if (oh, ow) == (h, w):
+        return x
+
+    def axis(n_in, n_out):
+        s = np.clip((np.arange(n_out, dtype=np.float64) + 0.5) * (n_in / n_out) - 0.5, 0, None)
+        i0 = np.minimum(np.floor(s).astype(np.int32), n_in - 1)
+        return mx.array(i0), mx.array(np.minimum(i0 + 1, n_in - 1)), mx.array((s - i0).astype(np.float32))
+
+    y0, y1, fy = axis(h, oh)
+    x0, x1, fx = axis(w, ow)
+    fy, fx = fy.astype(x.dtype)[None, :, None, None], fx.astype(x.dtype)[None, None, :, None]
+    top, bot = x[:, y0], x[:, y1]
+    top = top[:, :, x0] * (1 - fx) + top[:, :, x1] * fx
+    bot = bot[:, :, x0] * (1 - fx) + bot[:, :, x1] * fx
+    return top * (1 - fy) + bot * fy
+
+
+class LayerNorm2d(nn.Module):
+    """SAM's channel LayerNorm (eps 1e-6), statistics in float32."""
+
+    def __init__(self, c, eps=1e-6):
+        super().__init__()
+        self.weight, self.bias, self.eps = mx.ones((c,)), mx.zeros((c,)), eps
+
+    def __call__(self, x):
+        y = x.astype(mx.float32)
+        u = y.mean(-1, keepdims=True)
+        s = ((y - u) ** 2).mean(-1, keepdims=True)
+        return (self.weight * ((y - u) * mx.rsqrt(s + self.eps)).astype(x.dtype) + self.bias)
+
+
+
+def bicubic_matrix(n_in: int, n_out: int) -> np.ndarray:
+    """torch F.interpolate(mode="bicubic", align_corners=False) along one axis as an (n_out, n_in) matrix
+    (A = -0.75, source index not clamped, taps clamped to the border)."""
+    A = -0.75
+    m = np.zeros((n_out, n_in), np.float64)
+    scale = n_in / n_out
+    for o in range(n_out):
+        s = (o + 0.5) * scale - 0.5
+        f = int(np.floor(s))
+        t = s - f
+        w = [((A * (t + 1) - 5 * A) * (t + 1) + 8 * A) * (t + 1) - 4 * A,
+             ((A + 2) * t - (A + 3)) * t * t + 1,
+             ((A + 2) * (1 - t) - (A + 3)) * (1 - t) * (1 - t) + 1,
+             ((A * (2 - t) - 5 * A) * (2 - t) + 8 * A) * (2 - t) - 4 * A]
+        for j in range(4):
+            m[o, min(max(f - 1 + j, 0), n_in - 1)] += w[j]
+    return m.astype(np.float32)
+
+
+def bicubic(x, size):
+    b, h, w, c = x.shape
+    if (h, w) == tuple(size):
+        return x
+    mh, mw = mx.array(bicubic_matrix(h, size[0])), mx.array(bicubic_matrix(w, size[1]))
+    y = x.astype(mx.float32)
+    y = mx.einsum("oh,bhwc->bowc", mh, y)
+    y = mx.einsum("pw,bowc->bopc", mw, y)
+    return y.astype(x.dtype)
+
+
+def CL(cin, cout, k=1, s=1, g=1, bias=True):
+    """A ConvLayer with its BatchNorm folded in: a linear layer when it is a plain 1x1 convolution."""
+    if k == 1 and s == 1 and g == 1:
+        return nn.Linear(cin, cout, bias=bias)
+    return nn.Conv2d(cin, cout, k, stride=s, padding=k // 2, groups=g, bias=bias)
+
+
+act = nn.gelu_approx                                   # the repo's "gelu" is GELU(approximate="tanh")
+
+
+class ResBlock(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.conv1, self.conv2 = CL(c, c, 3), CL(c, c, 3)
+
+    def __call__(self, x):
+        return self.conv2(act(self.conv1(x)))
+
+
+class FusedMBConv(nn.Module):
+    def __init__(self, cin, cout, s, e):
+        super().__init__()
+        mid = round(cin * e)
+        self.spatial_conv, self.point_conv = CL(cin, mid, 3, s), CL(mid, cout)
+
+    def __call__(self, x):
+        return self.point_conv(act(self.spatial_conv(x)))
+
+
+class MBConv(nn.Module):
+    def __init__(self, cin, cout, s, e):
+        super().__init__()
+        mid = round(cin * e)
+        self.inverted_conv, self.depth_conv, self.point_conv = CL(cin, mid), CL(mid, mid, 3, s, mid), CL(mid, cout)
+
+    def __call__(self, x):
+        return self.point_conv(act(self.depth_conv(act(self.inverted_conv(x)))))
+
+
+class Aggreg(nn.Module):
+    def __init__(self, c, heads, k=5):
+        super().__init__()
+        self.dw = nn.Conv2d(c, c, k, padding=k // 2, groups=c, bias=False)
+        self.pw = nn.Conv2d(c, c, 1, groups=3 * heads, bias=False)
+
+    def __call__(self, x):
+        return self.pw(self.dw(x))
+
+
+class LiteMLA(nn.Module):
+    def __init__(self, c, dim=32):
+        super().__init__()
+        self.dim, heads = dim, c // dim
+        self.qkv = CL(c, 3 * c, bias=False)
+        self.aggreg = [Aggreg(3 * c, heads)]
+        self.proj = CL(2 * c, c)
+
+    def __call__(self, x):
+        b, h, w, _ = x.shape
+        qkv = self.qkv(x)
+        qkv = mx.concatenate([qkv] + [a(qkv) for a in self.aggreg], -1).astype(mx.float32)
+        d = self.dim
+        t = qkv.reshape(b, h * w, -1, 3 * d).transpose(0, 2, 1, 3)     # (B, heads, N, 3d)
+        q, k, v = nn.relu(t[..., :d]), nn.relu(t[..., d:2 * d]), t[..., 2 * d:]
+        v = mx.concatenate([v, mx.ones(v.shape[:-1] + (1,), mx.float32)], -1)
+        vk = v.transpose(0, 1, 3, 2) @ k                                   # (B, heads, d + 1, d)
+        out = q @ vk.transpose(0, 1, 3, 2)                                 # (B, heads, N, d + 1)
+        out = out[..., :-1] / (out[..., -1:] + 1e-15)
+        out = out.transpose(0, 2, 1, 3).reshape(b, h, w, -1).astype(x.dtype)
+        return self.proj(out)
+
+
+class EViTBlock(nn.Module):
+    def __init__(self, c, dim, e):
+        super().__init__()
+        self.context_module = LiteMLA(c, dim)
+        self.local_module = MBConv(c, c, 1, e)
+
+    def __call__(self, x):
+        x = x + self.context_module(x)
+        return x + self.local_module(x)
+
+
+class Res(nn.Module):
+    """ResidualBlock: x + main(x), or main(x) alone for the downsampling blocks."""
+
+    def __init__(self, main, shortcut=True):
+        super().__init__()
+        self.main, self.shortcut = main, shortcut
+
+    def __call__(self, x):
+        return x + self.main(x) if self.shortcut else self.main(x)
+
+
+class Backbone(nn.Module):
+    """EfficientViTLargeBackbone with the default blocks (res, fmb, fmb, mb, att) and expands (1, 4, 4, 4, 6)."""
+
+    def __init__(self, wl, dl, qkv_dim):
+        super().__init__()
+        blocks, expand = ["res", "fmb", "fmb", "mb", "att"], [1, 4, 4, 4, 6]
+        self.stem = CL(3, wl[0], 3, 2)
+        self.stages = [[Res(ResBlock(wl[0])) for _ in range(dl[0])]]
+        cin = wl[0]
+        for i in range(1, 5):
+            w, kind = wl[i], blocks[i]
+            down = FusedMBConv(cin, w, 2, expand[i] * 4) if kind == "fmb" else MBConv(cin, w, 2, expand[i] * 4)
+            stage = [Res(down, False)]
+            for _ in range(dl[i]):
+                if kind == "att":
+                    stage.append(EViTBlock(w, qkv_dim, expand[i]))
+                elif kind == "fmb":
+                    stage.append(Res(FusedMBConv(w, w, 1, expand[i])))
+                else:
+                    stage.append(Res(MBConv(w, w, 1, expand[i])))
+            self.stages.append(stage)
+            cin = w
+
+    def __call__(self, x):
+        x = act(self.stem(x))
+        feats = []
+        for stage in self.stages:
+            for blk in stage:
+                x = blk(x)
+            feats.append(x)
+        return feats                                                       # stage0 .. stage4
+
+
+class Neck(nn.Module):
+    def __init__(self, wl, depth, e, width=256):
+        super().__init__()
+        self.input_ops = [CL(wl[4], width), CL(wl[3], width), CL(wl[2], width)]   # stage4, stage3, stage2
+        self.middle = [Res(FusedMBConv(width, width, 1, e)) for _ in range(depth)]
+        self.output = CL(width, 256)
+
+    def __call__(self, feats, hw):
+        x = sum(bicubic(op(f), hw) for op, f in zip(self.input_ops, (feats[4], feats[3], feats[2])))
+        for blk in self.middle:
+            x = blk(x)
+        return self.output(x)
+
+
+class ImageEncoder(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.backbone = Backbone(c["width_list"], c["depth_list"], c["qkv_dim"])
+        self.neck = Neck(c["width_list"], c["neck_depth"], c["neck_expand"])
+        self.norm = LayerNorm2d(256, eps=1e-5)
+        self.emb = c["image_size"] // 16
+
+    def __call__(self, x):
+        return self.norm(self.neck(self.backbone(x), (self.emb, self.emb)))
+
+
+# ------------------------------------------------------------------------------------------------ SAM decoder
+
+class SAMAttention(nn.Module):
+    def __init__(self, dim, heads, down=1):
+        super().__init__()
+        inner = dim // down
+        self.heads = heads
+        self.q_proj, self.k_proj, self.v_proj = nn.Linear(dim, inner), nn.Linear(dim, inner), nn.Linear(dim, inner)
+        self.out_proj = nn.Linear(inner, dim)
+
+    def __call__(self, q, k, v):
+        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
+        b, nq, c = q.shape
+        sp = lambda t: t.reshape(b, t.shape[1], self.heads, c // self.heads).transpose(0, 2, 1, 3)
+        q, k, v = sp(q), sp(k), sp(v)
+        a = mx.softmax((q @ k.transpose(0, 1, 3, 2)) * (c // self.heads) ** -0.5, axis=-1, precise=True)
+        return self.out_proj((a @ v).transpose(0, 2, 1, 3).reshape(b, nq, c))
+
+
+class MLPBlock(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.lin1, self.lin2 = nn.Linear(dim, hidden), nn.Linear(hidden, dim)
+
+    def __call__(self, x):
+        return self.lin2(nn.relu(self.lin1(x)))
+
+
+class TwoWayBlock(nn.Module):
+    def __init__(self, dim, heads, mlp_dim, skip_pe):
+        super().__init__()
+        self.skip_pe = skip_pe
+        self.self_attn = SAMAttention(dim, heads)
+        self.cross_attn_token_to_image = SAMAttention(dim, heads, 2)
+        self.cross_attn_image_to_token = SAMAttention(dim, heads, 2)
+        self.mlp = MLPBlock(dim, mlp_dim)
+        self.norm1, self.norm2, self.norm3, self.norm4 = (nn.LayerNorm(dim) for _ in range(4))
+
+    def __call__(self, q, k, qpe, kpe):
+        if self.skip_pe:
+            q = self.self_attn(q, q, q)
+        else:
+            t = q + qpe
+            q = q + self.self_attn(t, t, q)
+        q = self.norm1(q)
+        q = self.norm2(q + self.cross_attn_token_to_image(q + qpe, k + kpe, k))
+        q = self.norm3(q + self.mlp(q))
+        k = self.norm4(k + self.cross_attn_image_to_token(k + kpe, q + qpe, q))
+        return q, k
+
+
+class TwoWayTransformer(nn.Module):
+    def __init__(self, dim=256, heads=8, mlp_dim=2048, depth=2):
+        super().__init__()
+        self.layers = [TwoWayBlock(dim, heads, mlp_dim, i == 0) for i in range(depth)]
+        self.final_attn_token_to_image = SAMAttention(dim, heads, 2)
+        self.norm_final_attn = nn.LayerNorm(dim)
+
+    def __call__(self, img, pe, tokens):
+        b, h, w, c = img.shape
+        k, kpe, q = img.reshape(b, h * w, c), pe.reshape(b, h * w, c), tokens
+        for layer in self.layers:
+            q, k = layer(q, k, tokens, kpe)
+        q = self.norm_final_attn(q + self.final_attn_token_to_image(q + tokens, k + kpe, k))
+        return q, k
+
+
+class MLP(nn.Module):
+    def __init__(self, din, dh, dout, n):
+        super().__init__()
+        dims = [din] + [dh] * (n - 1)
+        self.layers = [nn.Linear(a, b) for a, b in zip(dims, [dh] * (n - 1) + [dout])]
+
+    def __call__(self, x):
+        for i, l in enumerate(self.layers):
+            x = nn.relu(l(x)) if i < len(self.layers) - 1 else l(x)
+        return x
+
+
+class PromptEncoder(nn.Module):
+    def __init__(self, dim, emb_hw, img_size, mask_in):
+        super().__init__()
+        self.emb_hw, self.img_size = emb_hw, img_size
+        self.positional_encoding_gaussian_matrix = mx.zeros((2, dim // 2))
+        self.point_embeddings = [mx.zeros((1, dim)) for _ in range(4)]
+        self.not_a_point_embed = mx.zeros((1, dim))
+        self.no_mask_embed = mx.zeros((1, dim))
+        self.mask_conv1 = nn.Conv2d(1, mask_in // 4, 2, stride=2)
+        self.mask_norm1 = LayerNorm2d(mask_in // 4)
+        self.mask_conv2 = nn.Conv2d(mask_in // 4, mask_in, 2, stride=2)
+        self.mask_norm2 = LayerNorm2d(mask_in)
+        self.mask_conv3 = nn.Conv2d(mask_in, dim, 1)
+
+    def pe(self, coords):                                   # coords in [0, 1], float32
+        c = (2 * coords - 1) @ self.positional_encoding_gaussian_matrix.astype(mx.float32)
+        c = 2 * np.pi * c
+        return mx.concatenate([mx.sin(c), mx.cos(c)], -1)
+
+    def dense_pe(self):
+        h, w = self.emb_hw
+        ys = (mx.arange(h, dtype=mx.float32) + 0.5) / h
+        xs = (mx.arange(w, dtype=mx.float32) + 0.5) / w
+        grid = mx.stack([mx.broadcast_to(xs[None, :], (h, w)), mx.broadcast_to(ys[:, None], (h, w))], -1)
+        return self.pe(grid)[None]                          # (1, h, w, C)
+
+    def __call__(self, points, labels, box, mask):
+        dt = self.no_mask_embed.dtype
+        sparse = []
+        s = mx.array([self.img_size, self.img_size], dtype=mx.float32)
+        if points is not None:
+            pts, lab = points + 0.5, labels
+            if box is None:
+                pts = mx.concatenate([pts, mx.zeros((pts.shape[0], 1, 2))], 1)
+                lab = mx.concatenate([lab, -mx.ones((lab.shape[0], 1))], 1)
+            e = self.pe(pts / s)
+            l = lab[..., None]
+            e = mx.where(l == -1, 0.0, e)
+            e = e + (l == -1) * self.not_a_point_embed.astype(mx.float32)
+            e = e + (l == 0) * self.point_embeddings[0].astype(mx.float32)
+            e = e + (l == 1) * self.point_embeddings[1].astype(mx.float32)
+            sparse.append(e)
+        if box is not None:
+            e = self.pe((box + 0.5).reshape(-1, 2, 2) / s)
+            e = e + mx.concatenate([self.point_embeddings[2], self.point_embeddings[3]], 0).astype(mx.float32)[None]
+            sparse.append(e)
+        sparse = mx.concatenate(sparse, 1).astype(dt)
+        if mask is not None:
+            m = nn.gelu(self.mask_norm1(self.mask_conv1(mask)))
+            dense = self.mask_conv3(nn.gelu(self.mask_norm2(self.mask_conv2(m))))
+        else:
+            h, w = self.emb_hw
+            dense = mx.broadcast_to(self.no_mask_embed.reshape(1, 1, 1, -1), (sparse.shape[0], h, w, self.no_mask_embed.shape[-1]))
+        return sparse, dense
+
+
+class MaskDecoder(nn.Module):
+    def __init__(self, dim=256, n_mask=3):
+        super().__init__()
+        self.n = n_mask + 1
+        self.transformer = TwoWayTransformer(dim)
+        self.iou_token = mx.zeros((1, dim))
+        self.mask_tokens = mx.zeros((self.n, dim))
+        self.up1 = nn.ConvTranspose2d(dim, dim // 4, 2, stride=2)
+        self.up_norm = LayerNorm2d(dim // 4)
+        self.up2 = nn.ConvTranspose2d(dim // 4, dim // 8, 2, stride=2)
+        self.output_hypernetworks_mlps = [MLP(dim, dim, dim // 8, 3) for _ in range(self.n)]
+        self.iou_prediction_head = MLP(dim, 256, self.n, 3)
+
+    def __call__(self, img, pe, sparse, dense):
+        b = sparse.shape[0]
+        out = mx.concatenate([self.iou_token, self.mask_tokens], 0)
+        tokens = mx.concatenate([mx.broadcast_to(out[None], (b,) + out.shape), sparse], 1)
+        src = mx.broadcast_to(img, (b,) + img.shape[1:]) + dense
+        pos = mx.broadcast_to(pe, (b,) + pe.shape[1:]).astype(src.dtype)
+        hs, src = self.transformer(src, pos, tokens)
+        _, h, w, c = img.shape
+        up = nn.gelu(self.up2(nn.gelu(self.up_norm(self.up1(src.reshape(b, h, w, c))))))
+        hyper = mx.stack([self.output_hypernetworks_mlps[i](hs[:, 1 + i]) for i in range(self.n)], 1)
+        bb, uh, uw, uc = up.shape
+        masks = (hyper @ up.reshape(b, uh * uw, uc).transpose(0, 2, 1)).reshape(b, self.n, uh, uw)
+        return masks, self.iou_prediction_head(hs[:, 0])
+
+
+
+class Model(nn.Module):
+    def __init__(self, config: dict = CONFIG):
+        super().__init__()
+        self.size, self.encoder_size = config["image_size"], config["encoder_size"]
+        e = self.size // 16
+        self.image_encoder = ImageEncoder(config)
+        self.prompt_encoder = PromptEncoder(config["prompt_dim"], (e, e), self.size, config["mask_in_chans"])
+        self.mask_decoder = MaskDecoder(config["prompt_dim"])
+
+
+def load(folder) -> Model:
+    folder = Path(folder)
+    config = json.loads((folder / "config.json").read_text())
+    model = Model(config)
+    q = config.get("quantization")
+    if q:
+        nn.quantize(model, group_size=q["group_size"], bits=q["bits"], class_predicate=quantizable)
+    model.load_weights(str(folder / "model.safetensors"), strict=True)
+    mx.eval(model.parameters())
+    return model
+
+
+def dtype(model):
+    return model.mask_decoder.iou_token.dtype
+
+
+def embed(model: Model, rgb: np.ndarray) -> dict:
+    """EfficientViTSamPredictor.set_image: the encoder sees the image with its longest side at 512 (PIL bilinear),
+    normalized and padded; prompts and masks live in the 1024 frame (input_size)."""
+    from PIL import Image
+    h, w = rgb.shape[:2]
+    shape = lambda n: (int(h * n / max(h, w) + 0.5), int(w * n / max(h, w) + 0.5))
+    eh, ew = shape(model.encoder_size)
+    x = rgb if max(h, w) == model.encoder_size else np.asarray(Image.fromarray(rgb).resize((ew, eh), Image.BILINEAR))
+    x = (x.astype(np.float32) - PIXEL_MEAN) / PIXEL_STD
+    x = np.pad(x, ((0, model.encoder_size - eh), (0, model.encoder_size - ew), (0, 0)))
+    feats = model.image_encoder(mx.array(x[None]).astype(dtype(model)))
+    mx.eval(feats)
+    return {"features": feats, "original_size": (h, w), "input_size": shape(model.size)}
+
+
+def predict(model: Model, emb: dict, points=None, labels=None, box=None, multimask=True):
+    """SamPredictor.predict: prompts in original-image pixels -> (masks bool (N, H, W), scores (N,), low-res
+    logits (N, 256, 256))."""
+    h, w = emb["original_size"]
+    nh, nw = emb["input_size"]
+    sx, sy = nw / w, nh / h
+    pts = lab = bx = None
+    if points is not None:
+        pts = mx.array(np.asarray(points, np.float32)[None] * np.array([sx, sy], np.float32))
+        lab = mx.array(np.asarray(labels, np.float32)[None])
+    if box is not None:
+        bx = mx.array(np.asarray(box, np.float32).reshape(1, 4) * np.array([sx, sy, sx, sy], np.float32))
+    sparse, dense = model.prompt_encoder(pts, lab, bx, None)
+    pe = model.prompt_encoder.dense_pe()
+    logits, iou = model.mask_decoder(emb["features"], pe, sparse, dense.astype(sparse.dtype))
+    sl = slice(1, None) if multimask else slice(0, 1)
+    logits, iou = logits[:, sl].astype(mx.float32), iou[:, sl].astype(mx.float32)
+    low = logits
+    m = resize(logits.transpose(0, 2, 3, 1), (model.size, model.size))[:, :nh, :nw]
+    m = resize(m, (h, w)).transpose(0, 3, 1, 2)
+    return np.array(m[0] > 0), np.array(iou[0]), np.array(low[0]), np.array(m[0])
+
+
+
+# ------------------------------------------------------------------------------------------------ conversion
+
+def convert(source: Path) -> dict:
+    """runs/source/weights/efficientvit_sam_l2.pt (a plain state dict) -> MLX names: BatchNorms folded, 1x1 convs as
+    linear weights, the module tree renamed to this port's."""
+    import torch
+    sd = {k: v.double().numpy() for k, v in torch.load(Path(source) / "weights" / "efficientvit_sam_l2.pt",
+                                                          map_location="cpu", weights_only=True).items()
+          if v.is_floating_point()}
+    out = {}
+
+    def put(dst, conv, bn=None, linear=False):
+        w = sd[conv + ".weight"]
+        b = sd.get(conv + ".bias", np.zeros(w.shape[0]))
+        if bn is not None and bn + ".running_mean" in sd:
+            g = sd[bn + ".weight"] / np.sqrt(sd[bn + ".running_var"] + 1e-5)
+            w, b = w * g[:, None, None, None], (b - sd[bn + ".running_mean"]) * g + sd[bn + ".bias"]
+        out[dst + ".weight"] = (w[:, :, 0, 0] if linear else w.transpose(0, 2, 3, 1)).astype(np.float32)
+        if conv + ".bias" in sd or (bn is not None and bn + ".running_mean" in sd):
+            out[dst + ".bias"] = b.astype(np.float32)
+
+    def layer(dst, src, linear):                                      # a ConvLayer: src.conv (+ src.norm)
+        put(dst, src + ".conv", src + ".norm", linear)
+
+    bb = "image_encoder.backbone.stages"
+    layer("image_encoder.backbone.stem", f"{bb}.0.op_list.0", False)
+    st = lambda i, j: f"{bb}.{i}.op_list.{j}"
+    n0 = sum(1 for k in sd if k.startswith(f"{bb}.0.op_list.") and k.endswith("main.conv1.conv.weight"))
+    for j in range(n0):
+        for c in ("conv1", "conv2"):
+            layer(f"image_encoder.backbone.stages.0.{j}.main.{c}", f"{st(0, j + 1)}.main.{c}", False)
+    for i in range(1, 5):
+        j = 0
+        while any(k.startswith(st(i, j) + ".") for k in sd):
+            dst, src = f"image_encoder.backbone.stages.{i}.{j}", st(i, j)
+            if src + ".context_module.main.qkv.conv.weight" in sd:          # EfficientViTBlock
+                cm, lm = src + ".context_module.main", src + ".local_module.main"
+                layer(dst + ".context_module.qkv", cm + ".qkv", True)
+                put(dst + ".context_module.aggreg.0.dw", cm + ".aggreg.0.0")
+                put(dst + ".context_module.aggreg.0.pw", cm + ".aggreg.0.1")
+                layer(dst + ".context_module.proj", cm + ".proj", True)
+                layer(dst + ".local_module.inverted_conv", lm + ".inverted_conv", True)
+                layer(dst + ".local_module.depth_conv", lm + ".depth_conv", False)
+                layer(dst + ".local_module.point_conv", lm + ".point_conv", True)
+            elif src + ".main.spatial_conv.conv.weight" in sd:              # FusedMBConv
+                layer(dst + ".main.spatial_conv", src + ".main.spatial_conv", False)
+                layer(dst + ".main.point_conv", src + ".main.point_conv", True)
+            else:                                                           # MBConv
+                layer(dst + ".main.inverted_conv", src + ".main.inverted_conv", True)
+                layer(dst + ".main.depth_conv", src + ".main.depth_conv", False)
+                layer(dst + ".main.point_conv", src + ".main.point_conv", True)
+            j += 1
+    nk = "image_encoder.neck"
+    for i in range(3):
+        layer(f"{nk}.input_ops.{i}", f"{nk}.input_ops.{i}.op_list.0", True)
+    j = 0
+    while f"{nk}.middle.op_list.{j}.main.spatial_conv.conv.weight" in sd:
+        layer(f"{nk}.middle.{j}.main.spatial_conv", f"{nk}.middle.op_list.{j}.main.spatial_conv", False)
+        layer(f"{nk}.middle.{j}.main.point_conv", f"{nk}.middle.op_list.{j}.main.point_conv", True)
+        j += 1
+    layer(f"{nk}.output", f"{nk}.output_ops.0.op_list.0", True)
+    out["image_encoder.norm.weight"] = sd["image_encoder.norm.weight"].astype(np.float32)
+    out["image_encoder.norm.bias"] = sd["image_encoder.norm.bias"].astype(np.float32)
+    # SAM's prompt encoder and mask decoder, as in segment_anything
+    ren = {"prompt_encoder.pe_layer.": "prompt_encoder.",
+           "prompt_encoder.mask_downscaling.0.": "prompt_encoder.mask_conv1.",
+           "prompt_encoder.mask_downscaling.1.": "prompt_encoder.mask_norm1.",
+           "prompt_encoder.mask_downscaling.3.": "prompt_encoder.mask_conv2.",
+           "prompt_encoder.mask_downscaling.4.": "prompt_encoder.mask_norm2.",
+           "prompt_encoder.mask_downscaling.6.": "prompt_encoder.mask_conv3.",
+           "mask_decoder.output_upscaling.0.": "mask_decoder.up1.",
+           "mask_decoder.output_upscaling.1.": "mask_decoder.up_norm.",
+           "mask_decoder.output_upscaling.3.": "mask_decoder.up2."}
+    for k, v in sd.items():
+        if not k.startswith(("prompt_encoder.", "mask_decoder.")):
+            continue
+        n = k
+        for a, b in ren.items():
+            if n.startswith(a):
+                n = b + n[len(a):]
+        if n.endswith(".weight") and (n.startswith("prompt_encoder.point_embeddings") or
+                                      n.split(".")[-2] in ("not_a_point_embed", "no_mask_embed", "iou_token", "mask_tokens")):
+            n = n[:-len(".weight")]
+        if v.ndim == 4:
+            v = v.transpose(1, 2, 3, 0) if ".up1." in n or ".up2." in n else v.transpose(0, 2, 3, 1)
+        out[n] = v.astype(np.float32)
+    return out

fp16/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4b4b159c9f752c7ba72eeeb4aa099cf8575e064e01bfec0ee8576227172561f
+size 122661143