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