fp16 build of ORMBG

fp16/config.json

@@ -0,0 +1,5 @@
+{
+ "model_type": "ormbg",
+ "input_size": 1024,
+ "precision": "fp16"
+}

fp16/mlx_model.py

@@ -0,0 +1,200 @@
+"""ORMBG (Open Remove Background Model, an IS-Net from xuebinqin/DIS) in MLX: a port of schirrmacher/ormbg's
+`ormbg/models/ormbg.py` (Apache-2.0).
+
+    import mlx_model
+    model = mlx_model.load("fp16")                     # 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 and numpy). NHWC; each BatchNorm is folded into the conv before it (exact
+in fp32, done at conversion), so the weights are conv kernels and biases only.
+"""
+
+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": "ormbg", "input_size": 1024}
+PRECISIONS = ["bf16", "fp16"]       # all convolutions: MLX can't quantize them (model.toml)
+
+
+def resize_bilinear(x, size):
+    """torch F.interpolate(mode="bilinear", align_corners=False) for NHWC."""
+    b, h, w, c = x.shape
+    oh, ow = size
+    if (oh, ow) == (h, w):
+        return x
+
+    def axis(n_in, n_out):
+        s = (np.arange(n_out, dtype=np.float64) + 0.5) * (n_in / n_out) - 0.5
+        s = np.clip(s, 0, None)
+        i0 = np.minimum(np.floor(s).astype(np.int32), n_in - 1)
+        i1 = np.minimum(i0 + 1, n_in - 1)
+        f = (s - i0).astype(np.float32)
+        return mx.array(i0), mx.array(i1), mx.array(f)
+
+    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
+
+
+def maxpool2(x):
+    """MaxPool2d(2, 2, ceil_mode=True), NHWC."""
+    b, h, w, c = x.shape
+    ph, pw = h % 2, w % 2
+    if ph or pw:
+        x = mx.pad(x, [(0, 0), (0, ph), (0, pw), (0, 0)], constant_values=-float("inf"))
+        h, w = h + ph, w + pw
+    return x.reshape(b, h // 2, 2, w // 2, 2, c).max(axis=(2, 4))
+
+
+class REBNCONV(nn.Module):
+    def __init__(self, cin, cout, d=1):
+        super().__init__()
+        self.conv_s1 = nn.Conv2d(cin, cout, 3, padding=d, dilation=d)
+
+    def __call__(self, x):
+        return nn.relu(self.conv_s1(x))
+
+
+class RSU(nn.Module):
+    """RSU-L (L = depth: 7, 6, 5, 4): an encoder-decoder with pooling; dilation 2 at the bottom."""
+
+    def __init__(self, depth, cin, mid, cout):
+        super().__init__()
+        self.depth = depth
+        self.rebnconvin = REBNCONV(cin, cout)
+        setattr(self, "rebnconv1", REBNCONV(cout, mid))
+        for i in range(2, depth):
+            setattr(self, f"rebnconv{i}", REBNCONV(mid, mid))
+        setattr(self, f"rebnconv{depth}", REBNCONV(mid, mid, d=2))
+        for i in range(depth - 1, 1, -1):
+            setattr(self, f"rebnconv{i}d", REBNCONV(mid * 2, mid))
+        self.rebnconv1d = REBNCONV(mid * 2, cout)
+
+    def __call__(self, x):
+        hxin = self.rebnconvin(x)
+        enc, hx = [], hxin
+        for i in range(1, self.depth):
+            hx = getattr(self, f"rebnconv{i}")(hx)
+            enc.append(hx)
+            if i < self.depth - 1:
+                hx = maxpool2(hx)
+        hx = getattr(self, f"rebnconv{self.depth}")(hx)
+        for i in range(self.depth - 1, 0, -1):
+            hx = getattr(self, f"rebnconv{i}d")(mx.concatenate([hx, enc[i - 1]], -1))
+            if i > 1:
+                hx = resize_bilinear(hx, enc[i - 2].shape[1:3])
+        return hx + hxin
+
+
+class RSU4F(nn.Module):
+    """RSU-4F: no pooling, dilations 1, 2, 4, 8."""
+
+    def __init__(self, cin, mid, cout):
+        super().__init__()
+        self.rebnconvin = REBNCONV(cin, cout)
+        self.rebnconv1 = REBNCONV(cout, mid)
+        self.rebnconv2 = REBNCONV(mid, mid, 2)
+        self.rebnconv3 = REBNCONV(mid, mid, 4)
+        self.rebnconv4 = REBNCONV(mid, mid, 8)
+        self.rebnconv3d = REBNCONV(mid * 2, mid, 4)
+        self.rebnconv2d = REBNCONV(mid * 2, mid, 2)
+        self.rebnconv1d = REBNCONV(mid * 2, cout)
+
+    def __call__(self, x):
+        hxin = self.rebnconvin(x)
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+        hx4 = self.rebnconv4(hx3)
+        hx3d = self.rebnconv3d(mx.concatenate([hx4, hx3], -1))
+        hx2d = self.rebnconv2d(mx.concatenate([hx3d, hx2], -1))
+        return self.rebnconv1d(mx.concatenate([hx2d, hx1], -1)) + hxin
+
+
+class Model(nn.Module):
+    def __init__(self, config: dict = CONFIG):
+        super().__init__()
+        self.size = config["input_size"]
+        self.conv_in = nn.Conv2d(3, 64, 3, stride=2, padding=1)
+        self.stage1, self.stage2 = RSU(7, 64, 32, 64), RSU(6, 64, 32, 128)
+        self.stage3, self.stage4 = RSU(5, 128, 64, 256), RSU(4, 256, 128, 512)
+        self.stage5, self.stage6 = RSU4F(512, 256, 512), RSU4F(512, 256, 512)
+        self.stage5d, self.stage4d = RSU4F(1024, 256, 512), RSU(4, 1024, 128, 256)
+        self.stage3d, self.stage2d, self.stage1d = RSU(5, 512, 64, 128), RSU(6, 256, 32, 64), RSU(7, 128, 16, 64)
+        self.side1 = nn.Conv2d(64, 1, 3, padding=1)
+
+    def __call__(self, x):
+        """x: (B, S, S, 3) in [0, 1] -> (B, S, S, 1) foreground probability (side output 1, the one used)."""
+        hxin = self.conv_in(x)
+        hx1 = self.stage1(hxin)
+        hx2 = self.stage2(maxpool2(hx1))
+        hx3 = self.stage3(maxpool2(hx2))
+        hx4 = self.stage4(maxpool2(hx3))
+        hx5 = self.stage5(maxpool2(hx4))
+        hx6 = self.stage6(maxpool2(hx5))
+        up = lambda a, b: resize_bilinear(a, b.shape[1:3])
+        hx5d = self.stage5d(mx.concatenate([up(hx6, hx5), hx5], -1))
+        hx4d = self.stage4d(mx.concatenate([up(hx5d, hx4), hx4], -1))
+        hx3d = self.stage3d(mx.concatenate([up(hx4d, hx3), hx3], -1))
+        hx2d = self.stage2d(mx.concatenate([up(hx3d, hx2), hx2], -1))
+        hx1d = self.stage1d(mx.concatenate([up(hx2d, hx1), hx1], -1))
+        return mx.sigmoid(resize_bilinear(self.side1(hx1d), x.shape[1:3]))
+
+
+def load(folder) -> Model:
+    folder = Path(folder)
+    model = Model(json.loads((folder / "config.json").read_text()))
+    model.load_weights(str(folder / "model.safetensors"), strict=True)
+    mx.eval(model.parameters())
+    return model
+
+
+def predict(model: Model, rgb: np.ndarray) -> np.ndarray:
+    """The raw model output at the original size, float32 HxW in [0, 1] (as ormbg/inference.py, before min-max)."""
+    dt = model.conv_in.weight.dtype
+    h, w, _ = rgb.shape
+    x = mx.array(rgb.astype(np.float32))[None]
+    x = mx.floor(resize_bilinear(x, (model.size, model.size))) / 255.0      # the original rounds to uint8 first
+    y = model(x.astype(dt)).astype(mx.float32)
+    return np.array(resize_bilinear(y, (h, w)))[0, :, :, 0]
+
+
+def matte(model: Model, rgb: np.ndarray) -> np.ndarray:
+    """HxWx3 uint8 -> HxW uint8 alpha, min-max normalized like ormbg/inference.py."""
+    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/models/ormbg.pth -> conv weights with each BatchNorm folded in (fp32), MLX layout."""
+    import torch
+    sd = {k: v.double() for k, v in torch.load(Path(source) / "models" / "ormbg.pth", map_location="cpu",
+                                                  weights_only=True).items()}
+    out = {}
+    for k, w in sd.items():
+        if not k.endswith(".weight") or w.ndim != 4:
+            continue
+        base = k[:-len(".weight")]
+        b = sd.get(base + ".bias", torch.zeros(w.shape[0], dtype=torch.float64))
+        bn = base.replace("conv_s1", "bn_s1")
+        if bn != base and bn + ".running_mean" in sd:
+            g = sd[bn + ".weight"] / torch.sqrt(sd[bn + ".running_var"] + 1e-5)
+            w = w * g[:, None, None, None]
+            b = (b - sd[bn + ".running_mean"]) * g + sd[bn + ".bias"]
+        if base.startswith("side") and base != "side1":
+            continue                                    # side outputs 2-6 only matter for training
+        out[base + ".weight"] = w.float().numpy().transpose(0, 2, 3, 1)
+        out[base + ".bias"] = b.float().numpy()
+    return out

fp16/model.safetensors

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