File size: 4,843 Bytes
ec825b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """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
|