Spaces:
Running on Zero
Running on Zero
| """Restoration network used by the Space. | |
| Same architecture and the same parameter names as the checkpoint on the Hub, so the | |
| weights load unchanged. LayerNorm2d is written with plain tensor ops instead of a | |
| custom autograd Function, which keeps it portable. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| class LayerNorm2d(nn.Module): | |
| def __init__(self, channels, eps=1e-6): | |
| super().__init__() | |
| self.weight = nn.Parameter(torch.ones(channels)) | |
| self.bias = nn.Parameter(torch.zeros(channels)) | |
| 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): | |
| x1, x2 = x.chunk(2, dim=1) | |
| return x1 * x2 | |
| class Block(nn.Module): | |
| def __init__(self, c, dw_expand=2, ffn_expand=2): | |
| super().__init__() | |
| dw = c * dw_expand | |
| self.conv1 = nn.Conv2d(c, dw, 1, 1, 0, bias=True) | |
| self.conv2 = nn.Conv2d(dw, dw, 3, 1, 1, groups=dw, bias=True) | |
| self.conv3 = nn.Conv2d(dw // 2, c, 1, 1, 0, bias=True) | |
| self.sca = nn.Sequential( | |
| nn.AdaptiveAvgPool2d(1), | |
| nn.Conv2d(dw // 2, dw // 2, 1, 1, 0, bias=True), | |
| ) | |
| self.sg = SimpleGate() | |
| ffn = ffn_expand * c | |
| self.conv4 = nn.Conv2d(c, ffn, 1, 1, 0, bias=True) | |
| self.conv5 = nn.Conv2d(ffn // 2, c, 1, 1, 0, bias=True) | |
| 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.norm1(inp) | |
| x = self.conv1(x) | |
| x = self.conv2(x) | |
| x = self.sg(x) | |
| x = x * self.sca(x) | |
| x = self.conv3(x) | |
| y = inp + x * self.beta | |
| x = self.conv4(self.norm2(y)) | |
| x = self.sg(x) | |
| x = self.conv5(x) | |
| return y + x * self.gamma | |
| class RestoreNet(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), | |
| global_residual=True): | |
| super().__init__() | |
| self.global_residual = global_residual | |
| self.intro = nn.Conv2d(img_channel, width, 3, 1, 1, bias=True) | |
| self.ending = nn.Conv2d(width, img_channel, 3, 1, 1, bias=True) | |
| self.encoders = nn.ModuleList() | |
| self.decoders = nn.ModuleList() | |
| self.ups = nn.ModuleList() | |
| self.downs = nn.ModuleList() | |
| chan = width | |
| for num in enc_blk_nums: | |
| self.encoders.append(nn.Sequential(*[Block(chan) for _ in range(num)])) | |
| self.downs.append(nn.Conv2d(chan, 2 * chan, 2, 2)) | |
| chan *= 2 | |
| self.middle_blks = nn.Sequential(*[Block(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(*[Block(chan) for _ in range(num)])) | |
| self.padder_size = 2 ** len(self.encoders) | |
| def forward(self, inp): | |
| x = self.intro(inp) | |
| encs = [] | |
| for encoder, down in zip(self.encoders, self.downs): | |
| x = encoder(x) | |
| encs.append(x) | |
| x = down(x) | |
| x = self.middle_blks(x) | |
| for decoder, up, skip in zip(self.decoders, self.ups, encs[::-1]): | |
| x = up(x) | |
| x = x + skip | |
| x = decoder(x) | |
| x = self.ending(x) | |
| if self.global_residual: | |
| x = x + inp | |
| return x | |