Blame for 4bit/mlx_model.py

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