Spaces:
Sleeping
Sleeping
| """NAFNet (megvii-research, MIT) rebuilt in PyTorch to load the mlx-community NHWC safetensors.""" | |
| import json | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from safetensors.numpy import load_file | |
| class LayerNorm2d(nn.Module): | |
| def __init__(self, c, eps=1e-6): | |
| super().__init__() | |
| self.weight = nn.Parameter(torch.ones(c)) | |
| self.bias = nn.Parameter(torch.zeros(c)) | |
| self.eps = eps | |
| def forward(self, x): | |
| mu = x.mean(1, keepdim=True) | |
| var = (x - mu).pow(2).mean(1, keepdim=True) | |
| y = (x - mu) / torch.sqrt(var + self.eps) | |
| return self.weight.view(1, -1, 1, 1) * y + self.bias.view(1, -1, 1, 1) | |
| class SimpleGate(nn.Module): | |
| def forward(self, x): | |
| a, b = x.chunk(2, dim=1) | |
| return a * b | |
| class NAFBlock(nn.Module): | |
| def __init__(self, c): | |
| super().__init__() | |
| dw = c * 2 | |
| self.conv1 = nn.Conv2d(c, dw, 1) | |
| self.conv2 = nn.Conv2d(dw, dw, 3, padding=1, groups=dw) | |
| self.conv3 = nn.Conv2d(dw // 2, c, 1) | |
| # Simplified channel attention; global average pooling (the tile size equals the TLC window). | |
| self.sca = nn.Conv2d(dw // 2, dw // 2, 1) | |
| self.sg = SimpleGate() | |
| self.conv4 = nn.Conv2d(c, 2 * c, 1) | |
| self.conv5 = nn.Conv2d(c, c, 1) | |
| self.norm1 = LayerNorm2d(c) | |
| self.norm2 = LayerNorm2d(c) | |
| self.beta = nn.Parameter(torch.zeros(1, c, 1, 1)) | |
| self.gamma = nn.Parameter(torch.zeros(1, c, 1, 1)) | |
| def forward(self, inp): | |
| x = self.sg(self.conv2(self.conv1(self.norm1(inp)))) | |
| x = x * self.sca(x.mean(dim=(2, 3), keepdim=True)) | |
| y = inp + self.conv3(x) * self.beta | |
| x = self.conv5(self.sg(self.conv4(self.norm2(y)))) | |
| return y + x * self.gamma | |
| class NAFNet(nn.Module): | |
| def __init__(self, img_channel=3, width=64, middle_blk_num=1, enc_blk_nums=(1, 1, 1, 28), dec_blk_nums=(1, 1, 1, 1), **_): | |
| super().__init__() | |
| self.intro = nn.Conv2d(img_channel, width, 3, padding=1) | |
| self.ending = nn.Conv2d(width, img_channel, 3, padding=1) | |
| self.encoders, self.decoders, self.ups, self.downs = nn.ModuleList(), nn.ModuleList(), nn.ModuleList(), nn.ModuleList() | |
| chan = width | |
| for num in enc_blk_nums: | |
| self.encoders.append(nn.Sequential(*[NAFBlock(chan) for _ in range(num)])) | |
| self.downs.append(nn.Conv2d(chan, 2 * chan, 2, 2)) | |
| chan *= 2 | |
| self.middle_blks = nn.Sequential(*[NAFBlock(chan) for _ in range(middle_blk_num)]) | |
| for num in dec_blk_nums: | |
| self.ups.append(nn.Sequential(nn.Conv2d(chan, chan * 2, 1, bias=False), nn.PixelShuffle(2))) | |
| chan //= 2 | |
| self.decoders.append(nn.Sequential(*[NAFBlock(chan) for _ in range(num)])) | |
| def forward(self, inp): | |
| x = self.intro(inp) | |
| skips = [] | |
| for enc, down in zip(self.encoders, self.downs): | |
| x = enc(x) | |
| skips.append(x) | |
| x = down(x) | |
| x = self.middle_blks(x) | |
| for dec, up, skip in zip(self.decoders, self.ups, skips[::-1]): | |
| x = dec(up(x) + skip) | |
| return self.ending(x) + inp | |
| def load(weights, config=None): | |
| cfg = config or {"width": 64, "middle_blk_num": 1, "enc_blk_nums": [1, 1, 1, 28], "dec_blk_nums": [1, 1, 1, 1]} | |
| model = NAFNet(**cfg) | |
| raw = load_file(weights) | |
| state = {} | |
| target = model.state_dict() | |
| for key, value in raw.items(): | |
| name = key | |
| if key.startswith("ups.") and key.count(".") == 2: | |
| name = key.replace(".weight", ".0.weight") # ups.N is Sequential(conv, PixelShuffle) | |
| if value.ndim == 4 and name.endswith("weight"): | |
| value = np.transpose(value, (0, 3, 1, 2)) # MLX (O,kH,kW,I) -> PyTorch (O,I,kH,kW) | |
| elif name.endswith(("beta", "gamma")): | |
| value = np.transpose(value, (0, 3, 1, 2)) | |
| state[name] = torch.from_numpy(np.ascontiguousarray(value)) | |
| missing = set(target) - set(state) | |
| unused = set(state) - set(target) | |
| bad = [k for k in target if k in state and tuple(target[k].shape) != tuple(state[k].shape)] | |
| if missing or unused or bad: | |
| raise RuntimeError(f"missing={sorted(missing)[:8]} unused={sorted(unused)[:8]} shape={bad[:8]}") | |
| model.load_state_dict(state, strict=True) | |
| return model.eval() | |