bf16 build of Real-ESRGAN

bf16/config.json

@@ -0,0 +1,10 @@
+{
+ "model_type": "rrdbnet",
+ "num_in_ch": 3,
+ "num_out_ch": 3,
+ "num_feat": 64,
+ "num_block": 23,
+ "num_grow_ch": 32,
+ "scale": 4,
+ "precision": "bf16"
+}

bf16/mlx_model.py

@@ -0,0 +1,124 @@
+"""Real-ESRGAN x4plus (RRDBNet) in MLX: a port of basicsr's `RRDBNet` (xinntao/Real-ESRGAN, BSD-3-Clause).
+
+    import mlx_model
+    model = mlx_model.load("bf16")                       # a folder with config.json + model.safetensors
+    out = mlx_model.upscale(model, rgb)                  # rgb: HxWx3 uint8 numpy -> (4H)x(4W)x3 uint8
+
+Runs on its own next to its weights (only mlx and numpy). NHWC throughout; the weights are the original
+`params_ema` with each conv kernel transposed to MLX's (out, kh, kw, in).
+"""
+
+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": "rrdbnet", "num_in_ch": 3, "num_out_ch": 3, "num_feat": 64, "num_block": 23,
+          "num_grow_ch": 32, "scale": 4}
+# Every layer is a convolution, which MLX can't quantize: only bf16 / fp16 (lab/zoo/real-esrgan/model.toml).
+PRECISIONS = ["bf16", "fp16"]
+
+
+def lrelu(x):
+    return nn.leaky_relu(x, 0.2)
+
+
+class RDB(nn.Module):
+    def __init__(self, f: int, g: int):
+        super().__init__()
+        self.conv1 = nn.Conv2d(f, g, 3, padding=1)
+        self.conv2 = nn.Conv2d(f + g, g, 3, padding=1)
+        self.conv3 = nn.Conv2d(f + 2 * g, g, 3, padding=1)
+        self.conv4 = nn.Conv2d(f + 3 * g, g, 3, padding=1)
+        self.conv5 = nn.Conv2d(f + 4 * g, f, 3, padding=1)
+
+    def __call__(self, x):
+        x1 = lrelu(self.conv1(x))
+        x2 = lrelu(self.conv2(mx.concatenate([x, x1], -1)))
+        x3 = lrelu(self.conv3(mx.concatenate([x, x1, x2], -1)))
+        x4 = lrelu(self.conv4(mx.concatenate([x, x1, x2, x3], -1)))
+        x5 = self.conv5(mx.concatenate([x, x1, x2, x3, x4], -1))
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    def __init__(self, f: int, g: int):
+        super().__init__()
+        self.rdb1, self.rdb2, self.rdb3 = RDB(f, g), RDB(f, g), RDB(f, g)
+
+    def __call__(self, x):
+        return self.rdb3(self.rdb2(self.rdb1(x))) * 0.2 + x
+
+
+def up2(x):
+    """Nearest-neighbour x2 (NHWC)."""
+    b, h, w, c = x.shape
+    x = mx.broadcast_to(x[:, :, None, :, None, :], (b, h, 2, w, 2, c))
+    return x.reshape(b, h * 2, w * 2, c)
+
+
+class Model(nn.Module):
+    def __init__(self, config: dict = CONFIG):
+        super().__init__()
+        f, g, cin, cout = config["num_feat"], config["num_grow_ch"], config["num_in_ch"], config["num_out_ch"]
+        self.scale = config["scale"]
+        assert self.scale == 4, "x4plus only (x2 needs basicsr's pixel_unshuffle input)"
+        self.conv_first = nn.Conv2d(cin, f, 3, padding=1)
+        self.body = [RRDB(f, g) for _ in range(config["num_block"])]
+        self.conv_body = nn.Conv2d(f, f, 3, padding=1)
+        self.conv_up1 = nn.Conv2d(f, f, 3, padding=1)
+        self.conv_up2 = nn.Conv2d(f, f, 3, padding=1)
+        self.conv_hr = nn.Conv2d(f, f, 3, padding=1)
+        self.conv_last = nn.Conv2d(f, cout, 3, padding=1)
+
+    def __call__(self, x):
+        """x: (B, H, W, 3) in [0, 1] -> (B, 4H, 4W, 3)."""
+        feat = self.conv_first(x)
+        body = feat
+        for blk in self.body:
+            body = blk(body)
+        feat = feat + self.conv_body(body)
+        feat = lrelu(self.conv_up1(up2(feat)))
+        feat = lrelu(self.conv_up2(up2(feat)))
+        return self.conv_last(lrelu(self.conv_hr(feat)))
+
+
+def load(folder) -> Model:
+    folder = Path(folder)
+    config = json.loads((folder / "config.json").read_text())
+    model = Model(config)
+    model.load_weights(str(folder / "model.safetensors"), strict=True)
+    mx.eval(model.parameters())
+    return model
+
+
+def upscale(model: Model, rgb: np.ndarray, tile: int = 256, pad: int = 16) -> np.ndarray:
+    """HxWx3 uint8 -> (4H)x(4W)x3 uint8, in tiles of `tile` input pixels (overlap `pad`) to bound memory."""
+    dt = model.conv_first.weight.dtype
+    x = rgb.astype(np.float32) / 255.0
+    h, w, _ = x.shape
+    s = model.scale
+    out = np.zeros((h * s, w * s, 3), np.float32)
+    for y0 in range(0, h, tile):
+        for x0 in range(0, w, tile):
+            y1, x1 = min(y0 + tile, h), min(x0 + tile, w)
+            ya, xa, yb, xb = max(y0 - pad, 0), max(x0 - pad, 0), min(y1 + pad, h), min(x1 + pad, w)
+            t = mx.array(x[None, ya:yb, xa:xb]).astype(dt)
+            o = np.array(model(t).astype(mx.float32))[0]
+            out[y0 * s:y1 * s, x0 * s:x1 * s] = o[(y0 - ya) * s:(y1 - ya) * s, (x0 - xa) * s:(x1 - xa) * s]
+    return (np.clip(out, 0, 1) * 255.0).round().astype(np.uint8)
+
+
+def convert(source: Path) -> dict:
+    """runs/source/RealESRGAN_x4plus.pth (params_ema) -> MLX names and layouts, float32."""
+    import torch
+    sd = torch.load(Path(source) / "RealESRGAN_x4plus.pth", map_location="cpu", weights_only=True)["params_ema"]
+    out = {}
+    for k, v in sd.items():
+        a = v.float().numpy()
+        out[k] = a.transpose(0, 2, 3, 1) if a.ndim == 4 else a
+    return out

bf16/model.safetensors

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