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