text: # pipeline
  
  # generator
  # input -> Conv(k9n64s1)-> PReLU-> (Conv(k3n64s1)-> BN-> PReLU-> Conv(k3n64s1)-> BN) * 5->
  # Conv(k3n64s1)-> PReLU-> (Conv-> PixelShuffle-> PReLU)-> Conv(k3n64s1) -> output
  import numpy as np
  import torch
  import torch.nn as nn
  
  import math
  
  # ResidualBlock
  class ResidualBlock(nn.Module):
      def __init__(self, channels):
          super(ResidualBlock, self).__init__()
          self.Conv1 = nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=3, stride=1, padding=1)
          self.bn1 = nn.BatchNorm2d(channels)
          self.prelu = nn.PReLU()
          self.Conv2 = nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=3, stride=1, padding=1)
          self.bn2 = nn.BatchNorm2d(channels)
  
      def forward(self, x):
          residual = self.Conv1(x)
          residual = self.bn1(residual)
          residual = self.prelu(residual)
          residual = self.Conv2(residual)
          residual = self.bn2(residual)
          # 여기서 x와 residual의 크기가 달라 연산이 어렵다고 나온다.
          # 그 이유는 Conv층을 통과할 때 shape가 작아지기 때문이다.
          return x + residual
  
  # UpsampleBlock
  # (Conv-> PixelShuffle-> PReLU)
  class UpsampleBlock(nn.Module):
      def __init__(self, in_channels, up_scale):
          super(UpsampleBlock, self).__init__()
          self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2, kernel_size=3, padding=1)
          self.pixel_shuffle = nn.PixelShuffle(up_scale)
          self.prelu = nn.PReLU()
  
      def forward(self, x):
          x = self.conv(x)
          x = self.pixel_shuffle(x)
          x = self.prelu(x)
          return x
  
  
  # Generator
  class Generator(nn.Module):
      def __init__(self, scale_factor):
          upsample_block_num = int(math.log(scale_factor, 2))
  
          super(Generator, self).__init__()
          self.block1 = nn.Sequential(
              nn.Conv2d(3, 64, kernel_size=9, padding=4),
              nn.PReLU()
          )
          self.block2 = ResidualBlock(64)
          self.block3 = ResidualBlock(64)
          self.block4 = ResidualBlock(64)
          self.block5 = ResidualBlock(64)
          self.block6 = ResidualBlock(64)
          self.block7 = nn.Sequential(
              nn.Conv2d(64, 64, kernel_size=3, padding=1),
              nn.BatchNorm2d(64)
          )
          block8 = [UpsampleBlock(64, 2) for _ in range(upsample_block_num)]
          block8.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))
          self.block8 = nn.Sequential(*block8)
  
      def forward(self, x):
          block1 = self.block1(x)
          block2 = self.block2(block1)
          block3 = self.block3(block2)
          block4 = self.block4(block3)
          block5 = self.block5(block4)
          block6 = self.block6(block5)
          block7 = self.block7(block6)
          block8 = self.block8(block1 + block7)
          return (torch.tanh(block8) + 1) / 2
  
  # Discriminator
  # input -> Conv(k3n64s1)-> Leaky ReLU-> (Conv(k3n64s2)-> BN-> Leaky ReLU)->
  # (Conv(k3n128s1)-> BN-> Leaky ReLU)-> (Conv(k3n128s2)-> BN-> Leaky ReLU)->
  # (Conv(k3n256s1)-> BN-> Leaky ReLU)-> (Conv(k3n256s2)-> BN-> Leaky ReLU)->
  # (Conv(k3n512s1)-> BN-> Leaky ReLU)-> (Conv(k3n512s2)-> BN-> Leaky ReLU)->
  # Conv (1024)-> Leaky ReLU-> Conv (1)-> Sigmoid -> 1 or 0
  class Discriminator(nn.Module):
      def __init__(self):
          super(Discriminator, self).__init__()
          self.net = nn.Sequential(
              nn.Conv2d(3, 64, kernel_size=3, padding=1),
              nn.LeakyReLU(0.2),
  
              nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
              nn.BatchNorm2d(64),
              nn.LeakyReLU(0.2),
  
              nn.Conv2d(64, 128, kernel_size=3, padding=1),
              nn.BatchNorm2d(128),
              nn.LeakyReLU(0.2),
  
              nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1),
              nn.BatchNorm2d(128),
              nn.LeakyReLU(0.2),
  
              nn.Conv2d(128, 256, kernel_size=3, padding=1),
              nn.BatchNorm2d(256),
              nn.LeakyReLU(0.2),
  
              nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1),
              nn.BatchNorm2d(256),
              nn.LeakyReLU(0.2),
  
              nn.Conv2d(256, 512, kernel_size=3, padding=1),
              nn.BatchNorm2d(512),
              nn.LeakyReLU(0.2),
  
              nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
              nn.BatchNorm2d(512),
              nn.LeakyReLU(0.2),
  
              nn.AdaptiveAvgPool2d(1),
              nn.Conv2d(512, 1024, kernel_size=1),
              nn.LeakyReLU(0.2),
              nn.Conv2d(1024, 1, kernel_size=1)
          )
  
      def forward(self, x):
          batch_size = x.size(0)
          netx = self.net(x.clone())
          return torch.sigmoid(netx.view(batch_size))
  