bf16 build of BEN2 Base

bf16/config.json

@@ -0,0 +1,20 @@
+{
+ "model_type": "ben2",
+ "input_size": 1024,
+ "embed_dim": 128,
+ "depths": [
+  2,
+  2,
+  18,
+  2
+ ],
+ "num_heads": [
+  4,
+  8,
+  16,
+  32
+ ],
+ "window_size": 12,
+ "emb_dim": 128,
+ "precision": "bf16"
+}

bf16/mlx_model.py

@@ -0,0 +1,488 @@
+"""BEN2 Base (PramaLLC/BEN2, MIT) in MLX: a port of the Hub repo's `BEN2.py` (`BEN_Base`: Swin-B on four 512 px
+crops plus the whole image at 512 px, MCLM / MCRM cross-attention decoder, shallow full-resolution branch).
+
+    import mlx_model
+    model = mlx_model.load("bf16")                       # a folder with config.json + model.safetensors
+    alpha = mlx_model.matte(model, rgb)                  # rgb: HxWx3 uint8 -> HxW uint8 alpha (255 = foreground)
+
+Runs on its own next to its weights (mlx, numpy, pillow). NHWC throughout; InstanceNorm and the poolings are taken
+in float32. Quantized builds (8bit / 6bit / 4bit) quantize the linears (Swin, attention, feed-forward); the
+convolutions stay bf16. `refine_foreground` (the original's optional colour decontamination) isn't ported.
+"""
+
+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": "ben2", "input_size": 1024, "embed_dim": 128, "depths": [2, 2, 18, 2],
+          "num_heads": [4, 8, 16, 32], "window_size": 12, "emb_dim": 128}
+PRECISIONS = ["bf16", "fp16", "8bit", "6bit", "4bit"]
+MEAN, STD = np.array([0.485, 0.456, 0.406], np.float32), np.array([0.229, 0.224, 0.225], np.float32)
+GROUP = 64
+
+
+def quantizable(path: str, module) -> bool:
+    return isinstance(module, nn.Linear) and module.weight.shape[-1] % GROUP == 0
+
+
+# ------------------------------------------------------------------------------------------------ resize
+
+def resize(x, size, align_corners: bool):
+    """torch F.interpolate(mode="bilinear") for NHWC (no antialiasing)."""
+    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):
+        if align_corners:
+            s = np.arange(n_out, dtype=np.float64) * ((n_in - 1) / (n_out - 1) if n_out > 1 else 0.0)
+        else:
+            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)
+        i1 = np.minimum(i0 + 1, n_in - 1)
+        return mx.array(i0), mx.array(i1), mx.array((s - i0).astype(np.float32))
+
+    y0, y1, fy = axis(h, oh)
+    x0, x1, fx = axis(w, ow)
+    fy = fy.astype(x.dtype)[None, :, None, None]
+    fx = 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
+
+
+
+# ------------------------------------------------------------------------------------------------ Swin v1
+
+def window_partition(x, ws):
+    b, h, w, c = x.shape
+    x = x.reshape(b, h // ws, ws, w // ws, ws, c).transpose(0, 1, 3, 2, 4, 5)
+    return x.reshape(-1, ws * ws, c)
+
+
+def window_reverse(win, ws, h, w):
+    c = win.shape[-1]
+    b = win.shape[0] // ((h // ws) * (w // ws))
+    x = win.reshape(b, h // ws, w // ws, ws, ws, c).transpose(0, 1, 3, 2, 4, 5)
+    return x.reshape(b, h, w, c)
+
+
+def relative_index(ws: int) -> np.ndarray:
+    coords = np.stack(np.meshgrid(np.arange(ws), np.arange(ws), indexing="ij")).reshape(2, -1)
+    rel = (coords[:, :, None] - coords[:, None, :]).transpose(1, 2, 0) + (ws - 1)
+    return rel[:, :, 0] * (2 * ws - 1) + rel[:, :, 1]
+
+
+def shift_mask(hp: int, wp: int, ws: int, shift: int) -> np.ndarray:
+    img = np.zeros((1, hp, wp, 1), np.float32)
+    cnt = 0
+    for hs in (slice(0, -ws), slice(-ws, -shift), slice(-shift, None)):
+        for wsl in (slice(0, -ws), slice(-ws, -shift), slice(-shift, None)):
+            img[:, hs, wsl, :] = cnt
+            cnt += 1
+    m = img.reshape(1, hp // ws, ws, wp // ws, ws, 1).transpose(0, 1, 3, 2, 4, 5).reshape(-1, ws * ws)
+    m = m[:, None, :] - m[:, :, None]
+    return np.where(m != 0, -100.0, 0.0).astype(np.float32)
+
+
+class WindowAttention(nn.Module):
+    def __init__(self, dim, ws, heads):
+        super().__init__()
+        self.heads, self.scale = heads, (dim // heads) ** -0.5
+        self.relative_position_bias_table = mx.zeros(((2 * ws - 1) ** 2, heads))
+        self._index = mx.array(relative_index(ws).reshape(-1))
+        self.qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+    def __call__(self, x, mask=None):
+        b, n, c = x.shape
+        qkv = self.qkv(x).reshape(b, n, 3, self.heads, c // self.heads).transpose(2, 0, 3, 1, 4)
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        bias = self.relative_position_bias_table[self._index].reshape(n, n, -1).transpose(2, 0, 1)
+        attn = (q @ k.transpose(0, 1, 3, 2)) + bias[None].astype(q.dtype)
+        if mask is not None:
+            nw = mask.shape[0]
+            attn = attn.reshape(b // nw, nw, self.heads, n, n) + mask[None, :, None].astype(q.dtype)
+            attn = attn.reshape(b, self.heads, n, n)
+        attn = mx.softmax(attn, axis=-1, precise=True)
+        return self.proj((attn @ v).transpose(0, 2, 1, 3).reshape(b, n, c))
+
+
+class Mlp(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.fc1, self.fc2 = nn.Linear(dim, dim * 4), nn.Linear(dim * 4, dim)
+
+    def __call__(self, x):
+        return self.fc2(nn.gelu(self.fc1(x)))
+
+
+class SwinBlock(nn.Module):
+    def __init__(self, dim, heads, ws, shift):
+        super().__init__()
+        self.ws, self.shift = ws, shift
+        self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
+        self.attn = WindowAttention(dim, ws, heads)
+        self.mlp = Mlp(dim)
+
+    def __call__(self, x, h, w, mask):
+        b, L, c = x.shape
+        ws, s = self.ws, self.shift
+        y = self.norm1(x).reshape(b, h, w, c)
+        pr, pb = (ws - w % ws) % ws, (ws - h % ws) % ws
+        if pr or pb:
+            y = mx.pad(y, [(0, 0), (0, pb), (0, pr), (0, 0)])
+        hp, wp = h + pb, w + pr
+        if s:
+            y = mx.roll(y, (-s, -s), axis=(1, 2))
+        y = window_reverse(self.attn(window_partition(y, ws), mask if s else None), ws, hp, wp)
+        if s:
+            y = mx.roll(y, (s, s), axis=(1, 2))
+        y = y[:, :h, :w].reshape(b, h * w, c)
+        x = x + y
+        return x + self.mlp(self.norm2(x))
+
+
+class PatchMerging(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(4 * dim)
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+
+    def __call__(self, x, h, w):
+        b, L, c = x.shape
+        x = x.reshape(b, h, w, c)
+        if h % 2 or w % 2:
+            x = mx.pad(x, [(0, 0), (0, h % 2), (0, w % 2), (0, 0)])
+        x = mx.concatenate([x[:, 0::2, 0::2], x[:, 1::2, 0::2], x[:, 0::2, 1::2], x[:, 1::2, 1::2]], -1)
+        return self.reduction(self.norm(x.reshape(b, -1, 4 * c)))
+
+
+class SwinLayer(nn.Module):
+    def __init__(self, dim, depth, heads, ws, down):
+        super().__init__()
+        self.ws = ws
+        self.blocks = [SwinBlock(dim, heads, ws, 0 if i % 2 == 0 else ws // 2) for i in range(depth)]
+        if down:
+            self.downsample = PatchMerging(dim)
+        self._masks = {}
+
+    def __call__(self, x, h, w):
+        ws = self.ws
+        hp, wp = -(-h // ws) * ws, -(-w // ws) * ws
+        if (hp, wp) not in self._masks:
+            self._masks[(hp, wp)] = mx.array(shift_mask(hp, wp, ws, ws // 2))
+        mask = self._masks[(hp, wp)]
+        for blk in self.blocks:
+            x = blk(x, h, w, mask)
+        if "downsample" in self:
+            return x, self.downsample(x, h, w), (h + 1) // 2, (w + 1) // 2
+        return x, x, h, w
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.proj = nn.Conv2d(3, dim, 4, stride=4)
+        self.norm = nn.LayerNorm(dim)
+
+    def __call__(self, x):
+        b, h, w, _ = x.shape
+        if w % 4 or h % 4:
+            x = mx.pad(x, [(0, 0), (0, (4 - h % 4) % 4), (0, (4 - w % 4) % 4), (0, 0)])
+        return self.norm(self.proj(x))
+
+
+class Swin(nn.Module):
+    def __init__(self, embed_dim, depths, num_heads, window_size):
+        super().__init__()
+        self.patch_embed = PatchEmbed(embed_dim)
+        self.dims = [embed_dim * 2 ** i for i in range(4)]
+        self.layers = [SwinLayer(self.dims[i], depths[i], num_heads[i], window_size, i < 3) for i in range(4)]
+        for i, d in enumerate(self.dims):
+            setattr(self, f"norm{i}", nn.LayerNorm(d))
+
+    def __call__(self, x, emb=None):
+        x = self.patch_embed(x) if emb is None else emb                 # BEN2 also uses the embedding itself
+        b, h, w, c = x.shape
+        x = x.reshape(b, h * w, c)
+        outs = []
+        for i, layer in enumerate(self.layers):
+            xo, x, nh, nw = layer(x, h, w)
+            outs.append(getattr(self, f"norm{i}")(xo).reshape(b, h, w, self.dims[i]))
+            h, w = nh, nw
+        return outs
+
+
+# ------------------------------------------------------------------------------------------------ helpers
+
+def nearest(x, size):
+    """F.interpolate(mode="nearest") for NHWC: source index floor(dst * in / out)."""
+    b, h, w, c = x.shape
+    oh, ow = size
+    iy = mx.array(np.minimum((np.arange(oh) * (h / oh)).astype(np.int32), h - 1))
+    ix = mx.array(np.minimum((np.arange(ow) * (w / ow)).astype(np.int32), w - 1))
+    return x[:, iy][:, :, ix]
+
+
+def avg_pool_to(x, size):
+    """F.adaptive_avg_pool2d to a size that divides the input (always the case at 1024 px), in float32."""
+    b, h, w, c = x.shape
+    oh, ow = size
+    assert h % oh == 0 and w % ow == 0, (h, w, size)
+    y = x.astype(mx.float32).reshape(b, oh, h // oh, ow, w // ow, c).mean(axis=(2, 4))
+    return y.astype(x.dtype)
+
+
+def to_patches(x):
+    """'b c (hg h) (wg w) -> (hg wg b) c h w' with hg = wg = 2, NHWC."""
+    b, H, W, c = x.shape
+    x = x.reshape(b, 2, H // 2, 2, W // 2, c).transpose(1, 3, 0, 2, 4, 5)
+    return x.reshape(4 * b, H // 2, W // 2, c)
+
+
+def from_patches(x):
+    """'(hg wg b) c h w -> b c (hg h) (wg w)' with hg = wg = 2, NHWC."""
+    n, h, w, c = x.shape
+    b = n // 4
+    x = x.reshape(2, 2, b, h, w, c).transpose(2, 0, 3, 1, 4, 5)
+    return x.reshape(b, 2 * h, 2 * w, c)
+
+
+def sine_pos(h: int, w: int, d: int) -> np.ndarray:
+    """PositionEmbeddingSine(num_pos_feats = d // 2, normalize=True) as (h * w, d) tokens."""
+    n = d // 2
+    y = np.arange(1, h + 1, dtype=np.float32)[:, None].repeat(w, 1)
+    x = np.arange(1, w + 1, dtype=np.float32)[None, :].repeat(h, 0)
+    scale, eps = 2 * np.pi, 1e-6
+    y = (y - 0.5) / (y[-1:, :] + eps) * scale
+    x = (x - 0.5) / (x[:, -1:] + eps) * scale
+    dim_t = 10000 ** (2 * (np.arange(n, dtype=np.float32) // 2) / n)
+    px, py = x[:, :, None] / dim_t, y[:, :, None] / dim_t
+    px = np.stack((np.sin(px[:, :, 0::2]), np.cos(px[:, :, 1::2])), axis=3).reshape(h, w, -1)
+    py = np.stack((np.sin(py[:, :, 0::2]), np.cos(py[:, :, 1::2])), axis=3).reshape(h, w, -1)
+    return np.concatenate((py, px), axis=2).reshape(h * w, d).astype(np.float32)
+
+
+class InstanceNorm(nn.Module):
+    """nn.InstanceNorm2d(affine=False), statistics in float32."""
+
+    def __call__(self, x):
+        y = x.astype(mx.float32)
+        mu = y.mean(axis=(1, 2), keepdims=True)
+        var = ((y - mu) ** 2).mean(axis=(1, 2), keepdims=True)
+        return ((y - mu) * mx.rsqrt(var + 1e-5)).astype(x.dtype)
+
+
+class CBR(nn.Module):
+    """make_cbr / make_cbg: conv 3x3, InstanceNorm, GELU."""
+
+    def __init__(self, cin, cout):
+        super().__init__()
+        self.conv = nn.Conv2d(cin, cout, 3, padding=1)
+        self.norm = InstanceNorm()
+
+    def __call__(self, x):
+        return nn.gelu(self.norm(self.conv(x)))
+
+
+class MHA(nn.Module):
+    """nn.MultiheadAttention with one head, batch-first here: q (B, Lq, E), kv (B, Lk, E)."""
+
+    def __init__(self, d):
+        super().__init__()
+        self.d = d
+        self.in_proj = nn.Linear(d, 3 * d)
+        self.out_proj = nn.Linear(d, d)
+
+    def __call__(self, q, k, v):
+        w, b = self.in_proj.weight, self.in_proj.bias
+        if "scales" in self.in_proj:                          # quantized: run the whole projection, take slices
+            q = self.in_proj(q)[..., :self.d]
+            k = self.in_proj(k)[..., self.d:2 * self.d]
+            v = self.in_proj(v)[..., 2 * self.d:]
+        else:
+            d = self.d
+            q = q @ w[:d].T + b[:d]
+            k = k @ w[d:2 * d].T + b[d:2 * d]
+            v = v @ w[2 * d:].T + b[2 * d:]
+        a = mx.softmax((q * self.d ** -0.5) @ k.transpose(0, 2, 1), axis=-1, precise=True)
+        return self.out_proj(a @ v)
+
+
+class MCLM(nn.Module):
+    def __init__(self, d, pool_ratios=(1, 4, 8)):
+        super().__init__()
+        self.attention = [MHA(d) for _ in range(5)]
+        self.linear1, self.linear2 = nn.Linear(d, 2 * d), nn.Linear(2 * d, d)
+        self.linear3, self.linear4 = nn.Linear(d, 2 * d), nn.Linear(2 * d, d)
+        self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        self.pool_ratios, self.d = pool_ratios, d
+
+    def __call__(self, l, g):
+        """l (4, h, w, c) crops, g (1, h, w, c) whole image -> (5, h, w, c)."""
+        n, h, w, c = l.shape
+        cat = from_patches(l)                                              # (1, 2h, 2w, c)
+        pools, poses = [], []
+        for r in self.pool_ratios:
+            th, tw = round(h / r), round(w / r)
+            pools.append(avg_pool_to(cat, (th, tw)).reshape(1, th * tw, c))
+            poses.append(sine_pos(th, tw, c))
+        pools = mx.concatenate(pools, 1)
+        p_pos = mx.array(np.concatenate(poses, 0))[None].astype(l.dtype)
+        g_pos = mx.array(sine_pos(h, w, c))[None].astype(l.dtype)
+        gt = g.reshape(1, h * w, c)
+        gt = self.norm1(gt + self.attention[0](gt + g_pos, pools + p_pos, pools))
+        gt = self.norm2(gt + self.linear2(nn.gelu(self.linear1(gt))))
+        # each crop attends to its quarter of the whole-image tokens
+        gq = gt.reshape(2, h // 2, 2, w // 2, c).transpose(0, 2, 1, 3, 4).reshape(4, (h // 2) * (w // 2), c)
+        lt = l.reshape(4, h * w, c)
+        lt = lt + mx.concatenate([self.attention[i + 1](lt[i:i + 1], gq[i:i + 1], gq[i:i + 1]) for i in range(4)], 0)
+        lt = self.norm1(lt)
+        lt = self.norm2(lt + self.linear4(nn.gelu(self.linear3(lt))))
+        return mx.concatenate([lt, gt], 0).reshape(5, h, w, c)
+
+
+class MCRM(nn.Module):
+    def __init__(self, d, pool_ratios=(2, 4, 8)):
+        super().__init__()
+        self.attention = [MHA(d) for _ in range(4)]
+        self.linear3, self.linear4 = nn.Linear(d, 2 * d), nn.Linear(2 * d, d)
+        self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        self.sal_conv = nn.Conv2d(d, 1, 1)
+        self.pool_ratios = pool_ratios
+
+    def __call__(self, x):
+        n, h, w, c = x.shape
+        loc, glb = x[:4], x[4:]
+        pglb = to_patches(glb)                                             # (4, h/2, w/2, c)
+        tam = mx.sigmoid(self.sal_conv(glb))
+        loc = loc * to_patches(nearest(tam, (2 * h, 2 * w)))
+        pools = []
+        for r in self.pool_ratios:
+            th, tw = round(h / r), round(w / r)
+            pools.append(avg_pool_to(pglb, (th, tw)).reshape(4, th * tw, c))
+        pools = mx.concatenate(pools, 1)
+        lt = loc.reshape(4, h * w, c)
+        lt = lt + mx.concatenate([self.attention[i](lt[i:i + 1], pools[i:i + 1], pools[i:i + 1]) for i in range(4)], 0)
+        lt = self.norm1(lt)
+        lt = self.norm2(lt + self.linear4(nn.gelu(self.linear3(lt))))
+        src = lt.reshape(4, h, w, c)
+        glb = glb + nearest(from_patches(src), (h, w))
+        return mx.concatenate([src, glb], 0)
+
+
+class Head(nn.Module):
+    """insmask_head: conv, IN, GELU, conv, IN, GELU, conv."""
+
+    def __init__(self, d):
+        super().__init__()
+        self.c0, self.c3, self.c6 = nn.Conv2d(d, 384, 3, padding=1), nn.Conv2d(384, 384, 3, padding=1), nn.Conv2d(384, d, 3, padding=1)
+        self.norm = InstanceNorm()
+
+    def __call__(self, x):
+        x = nn.gelu(self.norm(self.c0(x)))
+        x = nn.gelu(self.norm(self.c3(x)))
+        return self.c6(x)
+
+
+class Model(nn.Module):
+    def __init__(self, config: dict = CONFIG):
+        super().__init__()
+        d = config["emb_dim"]
+        self.size = config["input_size"]
+        self.backbone = Swin(config["embed_dim"], config["depths"], config["num_heads"], config["window_size"])
+        dims = self.backbone.dims
+        self.output5, self.output4 = CBR(dims[3], d), CBR(dims[2], d)
+        self.output3, self.output2, self.output1 = CBR(dims[1], d), CBR(dims[0], d), CBR(config["embed_dim"], d)
+        self.multifieldcrossatt = MCLM(d)
+        self.conv1, self.conv2, self.conv3, self.conv4 = (CBR(d, d) for _ in range(4))
+        self.dec_blk1, self.dec_blk2, self.dec_blk3, self.dec_blk4 = (MCRM(d) for _ in range(4))
+        self.insmask_head = Head(d)
+        self.shallow = nn.Conv2d(3, d, 3, padding=1)
+        self.upsample1, self.upsample2 = CBR(d, d), CBR(d, d)
+        self.output = nn.Conv2d(d, 1, 3, padding=1)
+
+    def __call__(self, x):
+        """x: (1, S, S, 3) ImageNet-normalized -> (1, S, S, 1) foreground probability."""
+        b, h, w, _ = x.shape
+        shallow = self.shallow(x)
+        glb = resize(x, (h // 2, w // 2), False)
+        inp = mx.concatenate([to_patches(x), glb], 0)                     # 4 crops + the whole image
+        emb = self.backbone.patch_embed(inp)
+        f = [emb] + self.backbone(inp, emb)
+        e5, e4, e3 = self.output5(f[4]), self.output4(f[3]), self.output3(f[2])
+        e2, e1 = self.output2(f[1]), self.output1(f[0])
+        up = lambda a, ref: resize(a, ref.shape[1:3], False)
+        e5 = self.multifieldcrossatt(e5[:4], e5[4:])
+        e4 = self.conv4(self.dec_blk4(e4 + up(e5, e4)))
+        e3 = self.conv3(self.dec_blk3(e3 + up(e4, e3)))
+        e2 = self.conv2(self.dec_blk2(e2 + up(e3, e2)))
+        e1 = self.conv1(self.dec_blk1(e1 + up(e2, e1)))
+        out = from_patches(e1[:4])
+        out = out + up(e1[4:], out)
+        out = self.insmask_head(out)
+        out = out + up(shallow, out)
+        out = self.upsample1(nearest(out, (2 * out.shape[1], 2 * out.shape[2])))
+        out = out + up(shallow, out)
+        out = self.upsample2(nearest(out, (2 * out.shape[1], 2 * out.shape[2])))
+        return mx.sigmoid(self.output(out).astype(mx.float32))
+
+
+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 predict(model: Model, rgb: np.ndarray) -> np.ndarray:
+    """HxWx3 uint8 -> HxW float32 foreground probability, as BEN2.py before its min-max normalization (LANCZOS
+    resize to 1024, ImageNet normalization, bilinear resize of the result back)."""
+    from PIL import Image
+    h, w, _ = rgb.shape
+    s = model.size
+    x = np.asarray(Image.fromarray(rgb).resize((s, s), Image.LANCZOS), np.float32) / 255.0
+    x = mx.array(((x - MEAN) / STD)[None]).astype(model.shallow.weight.dtype)
+    y = model(x)
+    return np.array(resize(y, (h, w), False))[0, :, :, 0]
+
+
+def matte(model: Model, rgb: np.ndarray) -> np.ndarray:
+    """HxW uint8 alpha, min-max normalized like BEN2.py's postprocess_image."""
+    p = predict(model, rgb)
+    p = (p - p.min()) / max(p.max() - p.min(), 1e-8)
+    return (p * 255).astype(np.uint8)
+
+
+def convert(source: Path) -> dict:
+    """runs/source/model.safetensors -> MLX names and layouts (the side outputs, used in training only, dropped)."""
+    from safetensors.numpy import load_file
+    sd = load_file(str(Path(source) / "model.safetensors"))
+    ren = {"insmask_head.0": "insmask_head.c0", "insmask_head.3": "insmask_head.c3", "insmask_head.6": "insmask_head.c6",
+           "shallow.0": "shallow", "output.0": "output"}
+    out = {}
+    for k, v in sd.items():
+        if k.startswith("sideout") or k.endswith("relative_position_index"):
+            continue
+        n = k.replace("in_proj_weight", "in_proj.weight").replace("in_proj_bias", "in_proj.bias")
+        for a, b in ren.items():
+            if n.startswith(a + "."):
+                n = b + n[len(a):]
+        parts = n.split(".")
+        if len(parts) == 3 and parts[1] == "0" and parts[0][:-1] in ("output", "conv", "upsample"):
+            n = f"{parts[0]}.conv.{parts[2]}"                             # make_cbr Sequential -> CBR.conv
+        out[n] = v.transpose(0, 2, 3, 1) if v.ndim == 4 else v
+    return out

bf16/model.safetensors

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