text: import functools
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from CALAE.util.sr import initialize_weights, make_layer
  
  class ResidualBlock_noBN(nn.Module):
      '''Residual block w/o BN
      ---Conv-ReLU-Conv-+-
       |________________|
      '''
  
      def __init__(self, nf=64):
          super(ResidualBlock_noBN, self).__init__()
          self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
          self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
  
          # initialization
          initialize_weights([self.conv1, self.conv2], 0.1)
  
      def forward(self, x):
          identity = x
          out = F.relu(self.conv1(x), inplace=True)
          out = self.conv2(out)
          return identity + out
  
  class PA(nn.Module):
      '''PA is pixel attention'''
      def __init__(self, nf):
  
          super(PA, self).__init__()
          self.conv = nn.Conv2d(nf, nf, 1)
          self.sigmoid = nn.Sigmoid()
  
      def forward(self, x):
  
          y = self.conv(x)
          y = self.sigmoid(y)
          out = torch.mul(x, y)
  
          return out
      
  class ResidualBlock_noBN_PA(nn.Module):
      '''Residual block w/o BN
      ---Conv-ReLU-Conv-PA-+-
       |___________________|
      '''
  
      def __init__(self, nf=64):
          super(ResidualBlock_noBN_PA, self).__init__()
          self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
          self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
          self.pa = PA(nf)
          
          # initialization
          initialize_weights([self.conv1, self.conv2], 0.1)
  
      def forward(self, x):
          identity = x
          out = F.relu(self.conv1(x), inplace=True)
          out = self.conv2(out)
          out = self.pa(out)
          return identity + out
      
  class MSRResNet(nn.Module):
      ''' modified SRResNet'''
  
      def __init__(self, in_nc=3, out_nc=3, nf=64, nb=16, upscale=4):
          super(MSRResNet, self).__init__()
          self.upscale = upscale
  
          self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
          basic_block = functools.partial(ResidualBlock_noBN, nf=nf)
          self.recon_trunk = make_layer(basic_block, nb)
  
          # upsampling
          if self.upscale == 2:
              self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
              self.pixel_shuffle = nn.PixelShuffle(2)
          elif self.upscale == 3:
              self.upconv1 = nn.Conv2d(nf, nf * 9, 3, 1, 1, bias=True)
              self.pixel_shuffle = nn.PixelShuffle(3)
          elif self.upscale == 4:
              self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
              self.upconv2 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
              self.pixel_shuffle = nn.PixelShuffle(2)
  
          self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
          self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
  
          # activation function
          self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
  
          # initialization
          initialize_weights([self.conv_first, self.upconv1, self.HRconv, self.conv_last], 0.1)
          if self.upscale == 4:
              initialize_weights(self.upconv2, 0.1)
  
      def forward(self, x):
          fea = self.lrelu(self.conv_first(x))
          out = self.recon_trunk(fea)
  
          if self.upscale == 4:
              out = self.lrelu(self.pixel_shuffle(self.upconv1(out)))
              out = self.lrelu(self.pixel_shuffle(self.upconv2(out)))
          elif self.upscale == 3 or self.upscale == 2:
              out = self.lrelu(self.pixel_shuffle(self.upconv1(out)))
  
          out = self.conv_last(self.lrelu(self.HRconv(out)))
          base = F.interpolate(x, scale_factor=self.upscale, mode='bilinear', align_corners=False)
          out += base
          return out
      
  class MSRResNet_PA(nn.Module):
      ''' modified SRResNet + PA'''
  
      def __init__(self, in_nc=3, out_nc=3, nf=64, nb=16, upscale=4):
          super(MSRResNet_PA, self).__init__()
          self.upscale = upscale
  
          self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
          basic_block = functools.partial(ResidualBlock_noBN_PA, nf=nf)
          self.recon_trunk = make_layer(basic_block, nb)
  
          # upsampling
          if self.upscale == 2:
              self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
              self.pixel_shuffle = nn.PixelShuffle(2)
          elif self.upscale == 3:
              self.upconv1 = nn.Conv2d(nf, nf * 9, 3, 1, 1, bias=True)
              self.pixel_shuffle = nn.PixelShuffle(3)
          elif self.upscale == 4:
              self.upconv1 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
              self.upconv2 = nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
              self.pixel_shuffle = nn.PixelShuffle(2)
  
          self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
          self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
  
          # activation function
          self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
  
          # initialization
          initialize_weights([self.conv_first, self.upconv1, self.HRconv, self.conv_last],
                                       0.1)
          if self.upscale == 4:
              initialize_weights(self.upconv2, 0.1)
  
      def forward(self, x):
          fea = self.lrelu(self.conv_first(x))
          out = self.recon_trunk(fea)
  
          if self.upscale == 4:
              out = self.lrelu(self.pixel_shuffle(self.upconv1(out)))
              out = self.lrelu(self.pixel_shuffle(self.upconv2(out)))
          elif self.upscale == 3 or self.upscale == 2:
              out = self.lrelu(self.pixel_shuffle(self.upconv1(out)))
  
          out = self.conv_last(self.lrelu(self.HRconv(out)))
          base = F.interpolate(x, scale_factor=self.upscale, mode='bilinear', align_corners=False)
          out += base
          return out
  