| """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) |
|
|
|
|
| |
|
|
| def conv(cin, cout, k=1, s=1, p=0, g=1): |
| return nn.Conv2d(cin, cout, k, stride=s, padding=p, groups=g) |
|
|
|
|
| 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)))) |
|
|
|
|
| |
|
|
| 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): |
| 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] |
|
|
| 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"]) |
|
|
|
|
| |
|
|
| 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]) |
|
|
|
|
| |
|
|
| 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 |
|
|
