"""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