84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
|
||
|
|
from models.archs.RRDBNet_arch import RRDB
|
||
|
|
from models.archs.arch_util import make_layer, default_init_weights, ConvGnSilu, ConvGnLelu
|
||
|
|
from utils.util import checkpoint
|
||
|
|
|
||
|
|
|
||
|
|
class MultiLevelRRDB(nn.Module):
|
||
|
|
def __init__(self, nf, gc, levels):
|
||
|
|
super().__init__()
|
||
|
|
self.levels = levels
|
||
|
|
self.level_rrdbs = nn.ModuleList([RRDB(nf, growth_channels=gc) for i in range(levels)])
|
||
|
|
|
||
|
|
# Trunks should be fed in in order HR->LR
|
||
|
|
def forward(self, trunk):
|
||
|
|
for i in reversed(range(self.levels)):
|
||
|
|
lvl_scale = (2**i)
|
||
|
|
lvl_res = self.level_rrdbs[i](F.interpolate(trunk, scale_factor=1/lvl_scale, mode="area"), return_residual=True)
|
||
|
|
trunk = trunk + F.interpolate(lvl_res, scale_factor=lvl_scale, mode="nearest")
|
||
|
|
return trunk
|
||
|
|
|
||
|
|
|
||
|
|
class MultiResRRDBNet(nn.Module):
|
||
|
|
def __init__(self,
|
||
|
|
in_channels,
|
||
|
|
out_channels,
|
||
|
|
mid_channels=64,
|
||
|
|
l1_blocks=3,
|
||
|
|
l2_blocks=4,
|
||
|
|
l3_blocks=6,
|
||
|
|
growth_channels=32,
|
||
|
|
scale=4,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.scale = scale
|
||
|
|
self.in_channels = in_channels
|
||
|
|
|
||
|
|
self.conv_first = nn.Conv2d(in_channels, mid_channels, 7, stride=1, padding=3)
|
||
|
|
|
||
|
|
self.l3_blocks = nn.ModuleList([MultiLevelRRDB(mid_channels, growth_channels, 3) for _ in range(l1_blocks)])
|
||
|
|
self.l2_blocks = nn.ModuleList([MultiLevelRRDB(mid_channels, growth_channels, 2) for _ in range(l2_blocks)])
|
||
|
|
self.l1_blocks = nn.ModuleList([MultiLevelRRDB(mid_channels, growth_channels, 1) for _ in range(l3_blocks)])
|
||
|
|
self.block_levels = [self.l3_blocks, self.l2_blocks, self.l1_blocks]
|
||
|
|
|
||
|
|
self.conv_body = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
|
||
|
|
# upsample
|
||
|
|
self.conv_up1 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
|
||
|
|
self.conv_up2 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
|
||
|
|
self.conv_hr = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
|
||
|
|
self.conv_last = nn.Conv2d(mid_channels, out_channels, 3, 1, 1)
|
||
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
||
|
|
|
||
|
|
for m in [
|
||
|
|
self.conv_first, self.conv_first, self.conv_body, self.conv_up1,
|
||
|
|
self.conv_up2, self.conv_hr, self.conv_last
|
||
|
|
]:
|
||
|
|
if m is not None:
|
||
|
|
default_init_weights(m, 0.1)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
trunk = self.conv_first(x)
|
||
|
|
for block_set in self.block_levels:
|
||
|
|
for block in block_set:
|
||
|
|
trunk = checkpoint(block, trunk)
|
||
|
|
|
||
|
|
body_feat = self.conv_body(trunk)
|
||
|
|
feat = trunk + body_feat
|
||
|
|
|
||
|
|
# upsample
|
||
|
|
out = self.lrelu(
|
||
|
|
self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
|
||
|
|
if self.scale == 4:
|
||
|
|
out = self.lrelu(
|
||
|
|
self.conv_up2(F.interpolate(out, scale_factor=2, mode='nearest')))
|
||
|
|
else:
|
||
|
|
out = self.lrelu(self.conv_up2(out))
|
||
|
|
out = self.conv_last(self.lrelu(self.conv_hr(out)))
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
def visual_dbg(self, step, path):
|
||
|
|
pass
|