221 lines
7.9 KiB
Python
221 lines
7.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from models.arch_util import ConvGnLelu, default_init_weights, make_layer
|
|
from models.diffusion.nn import timestep_embedding
|
|
from trainer.networks import register_model
|
|
from utils.util import checkpoint
|
|
|
|
|
|
# Conditionally uses torch's checkpoint functionality if it is enabled in the opt file.
|
|
|
|
|
|
class ResidualDenseBlock(nn.Module):
|
|
"""Residual Dense Block.
|
|
|
|
Used in RRDB block in ESRGAN.
|
|
|
|
Args:
|
|
mid_channels (int): Channel number of intermediate features.
|
|
growth_channels (int): Channels for each growth.
|
|
"""
|
|
|
|
def __init__(self, mid_channels=64, growth_channels=32, embedding=False, init_weight=.1):
|
|
super(ResidualDenseBlock, self).__init__()
|
|
self.embedding = embedding
|
|
if embedding:
|
|
self.first_conv = ConvGnLelu(mid_channels, mid_channels, activation=True, norm=False, bias=True)
|
|
self.emb_layers = nn.Sequential(
|
|
nn.SiLU(),
|
|
nn.Linear(
|
|
mid_channels*4,
|
|
mid_channels,
|
|
),
|
|
)
|
|
for i in range(5):
|
|
out_channels = mid_channels if i == 4 else growth_channels
|
|
self.add_module(
|
|
f'conv{i + 1}',
|
|
nn.Conv2d(mid_channels + i * growth_channels, out_channels, 3,
|
|
1, 1))
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
|
for i in range(5):
|
|
default_init_weights(getattr(self, f'conv{i + 1}'), init_weight)
|
|
default_init_weights(self.conv5, 0)
|
|
|
|
self.normalize = nn.GroupNorm(num_groups=8, num_channels=mid_channels)
|
|
|
|
def forward(self, x, emb):
|
|
"""Forward function.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor with shape (n, c, h, w).
|
|
|
|
Returns:
|
|
Tensor: Forward results.
|
|
"""
|
|
if self.embedding:
|
|
x0 = self.first_conv(x)
|
|
emb_out = self.emb_layers(emb).type(x0.dtype)
|
|
while len(emb_out.shape) < len(x0.shape):
|
|
emb_out = emb_out[..., None]
|
|
x0 = x0 + emb_out
|
|
else:
|
|
x0 = x
|
|
x1 = self.lrelu(self.conv1(x0))
|
|
x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
|
|
x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
|
|
x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
|
|
x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
|
|
|
|
return self.normalize(x5 * .2 + x)
|
|
|
|
|
|
class RRDB(nn.Module):
|
|
"""Residual in Residual Dense Block.
|
|
|
|
Used in RRDB-Net in ESRGAN.
|
|
|
|
Args:
|
|
mid_channels (int): Channel number of intermediate features.
|
|
growth_channels (int): Channels for each growth.
|
|
"""
|
|
|
|
def __init__(self, mid_channels, growth_channels=32):
|
|
super(RRDB, self).__init__()
|
|
self.rdb1 = ResidualDenseBlock(mid_channels, growth_channels, embedding=True)
|
|
self.rdb2 = ResidualDenseBlock(mid_channels, growth_channels)
|
|
self.rdb3 = ResidualDenseBlock(mid_channels, growth_channels)
|
|
self.normalize = nn.GroupNorm(num_groups=8, num_channels=mid_channels)
|
|
self.residual_mult = nn.Parameter(torch.FloatTensor([.1]))
|
|
|
|
def forward(self, x, emb):
|
|
"""Forward function.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor with shape (n, c, h, w).
|
|
|
|
Returns:
|
|
Tensor: Forward results.
|
|
"""
|
|
out = self.rdb1(x, emb)
|
|
out = self.rdb2(out, emb)
|
|
out = self.rdb3(out, emb)
|
|
|
|
return self.normalize(out * self.residual_mult + x)
|
|
|
|
|
|
class RRDBNet(nn.Module):
|
|
"""Networks consisting of Residual in Residual Dense Block, which is used
|
|
in ESRGAN.
|
|
|
|
ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
|
|
Currently, it supports x4 upsampling scale factor.
|
|
|
|
Args:
|
|
in_channels (int): Channel number of inputs.
|
|
out_channels (int): Channel number of outputs.
|
|
mid_channels (int): Channel number of intermediate features.
|
|
Default: 64
|
|
num_blocks (int): Block number in the trunk network. Defaults: 23
|
|
growth_channels (int): Channels for each growth. Default: 32.
|
|
"""
|
|
|
|
def __init__(self,
|
|
in_channels,
|
|
out_channels,
|
|
mid_channels=64,
|
|
num_blocks=23,
|
|
growth_channels=32,
|
|
body_block=RRDB,
|
|
):
|
|
super(RRDBNet, self).__init__()
|
|
self.num_blocks = num_blocks
|
|
self.in_channels = in_channels
|
|
self.mid_channels = mid_channels
|
|
|
|
# The diffusion RRDB starts with a full resolution image and downsamples into a .25 working space
|
|
self.input_block = ConvGnLelu(in_channels, mid_channels, kernel_size=7, stride=1, activation=True, norm=False, bias=True)
|
|
self.down1 = ConvGnLelu(mid_channels, mid_channels, kernel_size=3, stride=2, activation=True, norm=False, bias=True)
|
|
self.down2 = ConvGnLelu(mid_channels, mid_channels, kernel_size=3, stride=2, activation=True, norm=False, bias=True)
|
|
|
|
# Guided diffusion uses a time embedding.
|
|
time_embed_dim = mid_channels * 4
|
|
self.time_embed = nn.Sequential(
|
|
nn.Linear(mid_channels, time_embed_dim),
|
|
nn.SiLU(),
|
|
nn.Linear(time_embed_dim, time_embed_dim),
|
|
)
|
|
|
|
self.body = make_layer(
|
|
body_block,
|
|
num_blocks,
|
|
mid_channels=mid_channels,
|
|
growth_channels=growth_channels)
|
|
|
|
self.conv_body = nn.Conv2d(self.mid_channels, self.mid_channels, 3, 1, 1)
|
|
# upsample
|
|
self.conv_up1 = nn.Conv2d(self.mid_channels, self.mid_channels, 3, 1, 1)
|
|
self.conv_up2 = nn.Conv2d(self.mid_channels*2, self.mid_channels, 3, 1, 1)
|
|
self.conv_up3 = None
|
|
self.conv_hr = nn.Conv2d(self.mid_channels*2, self.mid_channels, 3, 1, 1)
|
|
self.conv_last = nn.Conv2d(self.mid_channels, out_channels, 3, 1, 1)
|
|
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
|
self.normalize = nn.GroupNorm(num_groups=8, num_channels=self.mid_channels)
|
|
|
|
for m in [
|
|
self.conv_body, self.conv_up1,
|
|
self.conv_up2, self.conv_hr
|
|
]:
|
|
if m is not None:
|
|
default_init_weights(m, 1.0)
|
|
default_init_weights(self.conv_last, 0)
|
|
|
|
def forward(self, x, timesteps, low_res, correction_factors=None):
|
|
emb = self.time_embed(timestep_embedding(timesteps, self.mid_channels))
|
|
|
|
_, _, new_height, new_width = x.shape
|
|
upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
|
|
x = torch.cat([x, upsampled], dim=1)
|
|
|
|
if correction_factors is not None:
|
|
correction_factors = correction_factors.view(x.shape[0], -1, 1, 1).repeat(1, 1, new_height, new_width)
|
|
else:
|
|
correction_factors = torch.zeros((b, self.num_corruptions, new_height, new_width), dtype=torch.float, device=x.device)
|
|
x = torch.cat([x, correction_factors], dim=1)
|
|
|
|
d1 = self.input_block(x)
|
|
d2 = self.down1(d1)
|
|
feat = self.down2(d2)
|
|
for bl in self.body:
|
|
feat = checkpoint(bl, feat, emb)
|
|
feat = feat[:, :self.mid_channels]
|
|
feat = self.conv_body(feat)
|
|
|
|
# upsample
|
|
out = torch.cat([self.lrelu(
|
|
self.normalize(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))),
|
|
d2], dim=1)
|
|
out = torch.cat([self.lrelu(
|
|
self.normalize(self.conv_up2(F.interpolate(out, scale_factor=2, mode='nearest')))),
|
|
d1], dim=1)
|
|
out = self.conv_last(self.normalize(self.lrelu(self.conv_hr(out))))
|
|
|
|
return out
|
|
|
|
|
|
@register_model
|
|
def register_rrdb_diffusion(opt_net, opt):
|
|
return RRDBNet(**opt_net['args'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
model = RRDBNet(6,6)
|
|
x = torch.randn(1,3,128,128)
|
|
l = torch.randn(1,3,32,32)
|
|
t = torch.LongTensor([555])
|
|
y = model(x, t, l)
|
|
print(y.shape, y.mean(), y.std(), y.min(), y.max())
|