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(4): default_init_weights(getattr(self, f'conv{i + 1}'), init_weight) default_init_weights(self.conv5, 0) 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 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.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 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) 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=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) 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.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest'))), d2], dim=1) out = torch.cat([self.lrelu( self.conv_up2(F.interpolate(out, scale_factor=2, mode='nearest'))), d1], dim=1) out = self.conv_last(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())