fp16 build of MobileSAM v2

fp16/config.json

@@ -0,0 +1,31 @@
+{
+ "model_type": "mobile_sam",
+ "image_size": 1024,
+ "embed_dims": [
+  64,
+  128,
+  160,
+  320
+ ],
+ "depths": [
+  2,
+  2,
+  6,
+  2
+ ],
+ "num_heads": [
+  2,
+  4,
+  5,
+  10
+ ],
+ "window_sizes": [
+  7,
+  7,
+  14,
+  7
+ ],
+ "prompt_dim": 256,
+ "mask_in_chans": 16,
+ "precision": "fp16"
+}

fp16/mlx_model.py

@@ -0,0 +1,465 @@
+"""MobileSAM (ChaoningZhang/MobileSAM, Apache-2.0) in MLX: a port of `mobile_sam` (`build_sam_vit_t`: TinyViT-5M
+image encoder + SAM's prompt encoder and mask decoder), the segment-anything half of MobileSAM v2.
+
+    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, logits = mlx_model.predict(model, emb, points=[[x, y]], labels=[1], box=None, multimask=True)
+
+`predict` follows SamPredictor.predict: point / box prompts in pixel coordinates of the original image, masks at
+the original size (logits > 0), the IoU head's scores, and 256x256 low-res logits. Runs on its own next to its
+weights (mlx, numpy, pillow). NHWC; every BatchNorm folded into its conv at conversion.
+"""
+
+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": "mobile_sam", "image_size": 1024, "embed_dims": [64, 128, 160, 320], "depths": [2, 2, 6, 2],
+          "num_heads": [2, 4, 5, 10], "window_sizes": [7, 7, 14, 7], "prompt_dim": 256, "mask_in_chans": 16}
+PRECISIONS = ["bf16", "fp16", "8bit", "6bit", "4bit"]
+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)
+
+
+# ------------------------------------------------------------------------------------------------ TinyViT
+
+def conv(cin, cout, k=1, s=1, p=0, g=1):
+    return nn.Conv2d(cin, cout, k, stride=s, padding=p, groups=g)          # Conv2d_BN, BatchNorm folded
+
+
+class MBConv(nn.Module):
+    def __init__(self, c, ratio=4.0):
+        super().__init__()
+        h = int(c * ratio)
+        self.conv1, self.conv2, self.conv3 = conv(c, h), conv(h, h, 3, 1, 1, h), conv(h, c)
+
+    def __call__(self, x):
+        y = nn.gelu(self.conv1(x))
+        y = nn.gelu(self.conv2(y))
+        return nn.gelu(self.conv3(y) + x)
+
+
+class PatchMerging(nn.Module):
+    def __init__(self, dim, out):
+        super().__init__()
+        s = 1 if out in (320, 448, 576) else 2
+        self.conv1, self.conv2, self.conv3 = conv(dim, out), conv(out, out, 3, s, 1, out), conv(out, out)
+
+    def __call__(self, x):
+        return self.conv3(nn.gelu(self.conv2(nn.gelu(self.conv1(x)))))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, key_dim, heads, res):
+        super().__init__()
+        self.heads, self.key_dim, self.scale = heads, key_dim, key_dim ** -0.5
+        self.norm = nn.LayerNorm(dim)
+        self.qkv = nn.Linear(dim, 3 * key_dim * heads)
+        self.proj = nn.Linear(key_dim * heads, dim)
+        pts = [(i, j) for i in range(res) for j in range(res)]
+        offs, idx = {}, []
+        for p1 in pts:
+            for p2 in pts:
+                o = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
+                offs.setdefault(o, len(offs))
+                idx.append(offs[o])
+        self.attention_biases = mx.zeros((heads, len(offs)))
+        self._idx = mx.array(np.array(idx, np.int32).reshape(len(pts), len(pts)))
+
+    def __call__(self, x):
+        b, n, _ = x.shape
+        qkv = self.qkv(self.norm(x)).reshape(b, n, self.heads, 3 * self.key_dim)
+        q, k, v = mx.split(qkv, [self.key_dim, 2 * self.key_dim], axis=-1)
+        q, k, v = (t.transpose(0, 2, 1, 3) for t in (q, k, v))
+        a = (q @ k.transpose(0, 1, 3, 2)) * self.scale + self.attention_biases[:, self._idx].astype(q.dtype)
+        a = mx.softmax(a, axis=-1, precise=True)
+        return self.proj((a @ v).transpose(0, 2, 1, 3).reshape(b, n, -1))
+
+
+class Mlp(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.norm, self.fc1, self.fc2 = nn.LayerNorm(dim), nn.Linear(dim, hidden), nn.Linear(hidden, dim)
+
+    def __call__(self, x):
+        return self.fc2(nn.gelu(self.fc1(self.norm(x))))
+
+
+class TinyViTBlock(nn.Module):
+    def __init__(self, dim, heads, ws):
+        super().__init__()
+        self.ws = ws
+        self.attn = Attention(dim, dim // heads, heads, ws)
+        self.mlp = Mlp(dim, dim * 4)
+        self.local_conv = conv(dim, dim, 3, 1, 1, dim)
+
+    def __call__(self, x):
+        """x: (B, H, W, C)."""
+        b, h, w, c = x.shape
+        ws = self.ws
+        if h == ws and w == ws:
+            y = self.attn(x.reshape(b, h * w, c)).reshape(b, h, w, c)
+        else:
+            pb, pr = (ws - h % ws) % ws, (ws - w % ws) % ws
+            y = mx.pad(x, [(0, 0), (0, pb), (0, pr), (0, 0)]) if pb or pr else x
+            ph, pw = h + pb, w + pr
+            y = y.reshape(b, ph // ws, ws, pw // ws, ws, c).transpose(0, 1, 3, 2, 4, 5).reshape(-1, ws * ws, c)
+            y = self.attn(y)
+            y = y.reshape(b, ph // ws, pw // ws, ws, ws, c).transpose(0, 1, 3, 2, 4, 5).reshape(b, ph, pw, c)
+            y = y[:, :h, :w]
+        x = self.local_conv(x + y)
+        return x + self.mlp(x.reshape(b, h * w, c)).reshape(b, h, w, c)
+
+
+class Layer(nn.Module):
+    def __init__(self, i, dims, depth, heads, ws):
+        super().__init__()
+        d = dims[i]
+        self.blocks = [MBConv(d) for _ in range(depth)] if i == 0 else [TinyViTBlock(d, heads, ws) for _ in range(depth)]
+        if i < len(dims) - 1:
+            self.downsample = PatchMerging(d, dims[i + 1])
+
+    def __call__(self, x):
+        for b in self.blocks:
+            x = b(x)
+        return self.downsample(x) if "downsample" in self else x
+
+
+class TinyViT(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        d = c["embed_dims"]
+        self.patch_embed = [conv(3, d[0] // 2, 3, 2, 1), conv(d[0] // 2, d[0], 3, 2, 1)]
+        self.layers = [Layer(i, d, c["depths"][i], c["num_heads"][i], c["window_sizes"][i]) for i in range(4)]
+        self.neck_conv1 = nn.Conv2d(d[-1], 256, 1, bias=False)
+        self.neck_norm1 = LayerNorm2d(256)
+        self.neck_conv2 = nn.Conv2d(256, 256, 3, padding=1, bias=False)
+        self.neck_norm2 = LayerNorm2d(256)
+
+    def __call__(self, x):
+        x = self.patch_embed[1](nn.gelu(self.patch_embed[0](x)))
+        for layer in self.layers:
+            x = layer(x)
+        return self.neck_norm2(self.neck_conv2(self.neck_norm1(self.neck_conv1(x))))
+
+
+# ------------------------------------------------------------------------------------------------ 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 = config["image_size"]
+        e = self.size // 16
+        self.image_encoder = TinyViT(config)
+        self.prompt_encoder = PromptEncoder(config["prompt_dim"], (e, e), self.size, config["mask_in_chans"])
+        self.mask_decoder = MaskDecoder(config["prompt_dim"])
+
+
+# ------------------------------------------------------------------------------------------------ predictor
+
+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:
+    """SamPredictor.set_image: longest side to 1024 (PIL bilinear), normalize, pad bottom-right, encode."""
+    from PIL import Image
+    h, w = rgb.shape[:2]
+    s = model.size / max(h, w)
+    nh, nw = int(h * s + 0.5), int(w * s + 0.5)
+    x = np.asarray(Image.fromarray(rgb).resize((nw, nh), Image.BILINEAR), np.float32)
+    x = (x - PIXEL_MEAN) / PIXEL_STD
+    x = np.pad(x, ((0, model.size - nh), (0, model.size - nw), (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": (nh, nw)}
+
+
+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 _fold(sd, pre):
+    """Conv2d_BN (`pre`.c + `pre`.bn) -> one conv's (weight, bias) in MLX layout."""
+    w = sd[pre + ".c.weight"].astype(np.float64)
+    g = sd[pre + ".bn.weight"] / np.sqrt(sd[pre + ".bn.running_var"].astype(np.float64) + 1e-5)
+    b = sd[pre + ".bn.bias"] - sd[pre + ".bn.running_mean"] * g
+    return (w * g[:, None, None, None]).astype(np.float32).transpose(0, 2, 3, 1), b.astype(np.float32)
+
+
+def convert(source: Path) -> dict:
+    """runs/source/weights/mobile_sam.pt (a plain state dict) -> MLX names: BatchNorms folded, Sequentials renamed,
+    the classification head (norm_head / head) dropped."""
+    import torch
+    sd = {k: v.float().numpy() for k, v in torch.load(Path(source) / "weights" / "mobile_sam.pt", map_location="cpu",
+                                                         weights_only=True).items() if v.is_floating_point()}
+    out = {}
+    folded = sorted({k[:-len(".c.weight")] for k in sd if k.endswith(".c.weight")})
+    for pre in folded:
+        dst = pre.replace("patch_embed.seq.0", "patch_embed.0").replace("patch_embed.seq.2", "patch_embed.1")
+        out[dst + ".weight"], out[dst + ".bias"] = _fold(sd, pre)
+    ren = {"image_encoder.neck.0.": "image_encoder.neck_conv1.", "image_encoder.neck.1.": "image_encoder.neck_norm1.",
+           "image_encoder.neck.2.": "image_encoder.neck_conv2.", "image_encoder.neck.3.": "image_encoder.neck_norm2.",
+           "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 any(k.startswith(p + ".") for p in folded) or ".head." in k or ".norm_head." in k:
+            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
+    return out

fp16/model.safetensors

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