text: '''FastGAN model for Tensorflow.
  
  Reference:
    - Bingchen Liu, Yizhe Zhu, Kunpeng Song and Ahmed Elgammal. 
      [Towards Faster and Stabilized GAN Training for High-fidelity Few-shot Image Synthesis](
        https://arxiv.org/abs/2101.04775) 
    
  Author: Emilio Morales (mil.mor.mor@gmail.com)
          Jan 2022
  '''
  import tensorflow as tf
  from tensorflow.keras import layers
  from tensorflow_addons.layers import AdaptiveAveragePooling2D
  
  
  def normalize_2nd_moment(x, axis=1, eps=1e-8):
      return x * tf.math.rsqrt(tf.reduce_mean(tf.square(x), axis=axis, keepdims=True) + eps)
  
  
  class GLU(layers.Layer):
      def __init__(self, **kwargs):
          super(GLU, self).__init__(**kwargs)
  
      def call(self, x):
          self.nc = x.shape[-1] // 2
          if len(x.shape) == 2:
              return x[:, :self.nc] * tf.nn.sigmoid(x[:, self.nc:])
  
          elif len(x.shape) == 4:
              return x[:, :, :, :self.nc] * tf.nn.sigmoid(x[:, :, :, self.nc:])
  
  
  class NoiseInjection(layers.Layer):
      def __init__(self, **kwargs):
          super(NoiseInjection, self).__init__(**kwargs)
          self.weight = self.add_weight(
              'weight', shape=(1), initializer='zeros', 
              dtype=self.dtype, trainable=True)
  
      def call(self, feat):
          batch, height, width, _ = feat.shape
          noise = tf.random.normal((batch, height, width, 1), dtype=self.dtype)
          return feat + self.weight * noise
  
  
  class UpBlockComp(layers.Layer):
      def __init__(self, filters, kernel_size=3, initializer='orthogonal'):
          super(UpBlockComp, self).__init__()
          self.up_conv = tf.keras.Sequential([
              layers.UpSampling2D(2),
              layers.Conv2D(filters*2, kernel_size=kernel_size, 
                  padding='same', use_bias=False, kernel_initializer=initializer),
          ])
          self.noise_1 = NoiseInjection()
          self.act_1 = tf.keras.Sequential([
              layers.BatchNormalization(),
              GLU()
          ])
          self.conv = layers.Conv2D(filters*2, kernel_size=kernel_size, 
              padding='same', use_bias=False, kernel_initializer=initializer)
          self.noise_2 = NoiseInjection()
          self.act_2 = tf.keras.Sequential([
              layers.BatchNormalization(),
              GLU()
          ])
  
      def call(self, x):
          x = self.up_conv(x)
          x = self.noise_1(x)
          x = self.act_1(x)
          x = self.conv(x)
          x = self.noise_2(x)
          x = self.act_2(x)
          return x
      
  
  def upBlock(filters, kernel_size=3, initializer='orthogonal'):
      block = tf.keras.Sequential([
              layers.UpSampling2D(2),
              layers.Conv2D(filters*2, kernel_size=kernel_size, 
                  padding='same', use_bias=False, kernel_initializer=initializer),
              layers.BatchNormalization(),
              GLU()
      ])
      return block
  
  
  def SEBlock(filters, kernel_size=4, activation='swish', initializer='orthogonal'):
      block = tf.keras.Sequential([
                  AdaptiveAveragePooling2D(4),
                  layers.Conv2D(filters=filters, 
                      kernel_size=4, 
                      activation=activation, 
                      use_bias=False, kernel_initializer=initializer),
                  layers.Conv2D(filters=filters, 
                      kernel_size=1, 
                      activation='sigmoid', 
                      use_bias=False, kernel_initializer=initializer)
      ])
      return block
  
  
  class InitLayer(layers.Layer):
      def __init__(self, units=256, initializer='orthogonal'):
          super(InitLayer, self).__init__()
          self.conv = tf.keras.Sequential([
              layers.Conv2DTranspose(units * 2, 
                  kernel_size=4, 
                  use_bias=False, kernel_initializer=initializer),
                  layers.BatchNormalization(),
                  GLU()
          ])
          
      def call(self, x):
          x = tf.expand_dims(x, 1)
          x = tf.expand_dims(x, 1)
          return self.conv(x)
  
  
  class Generator(tf.keras.models.Model):
      def __init__(self, filters=1024, initializer='orthogonal'):
          super(Generator, self).__init__()
          self.init = InitLayer()
          
          self.up_8 = UpBlockComp(filters, initializer=initializer)
          self.up_16 = upBlock(filters // 2, initializer=initializer)
          self.up_32 = UpBlockComp(filters // 4, initializer=initializer)
         
          self.up_64 = upBlock(filters // 8, initializer=initializer)
          self.se_64 = SEBlock(filters // 8, initializer=initializer)
          
          self.up_128 = UpBlockComp(filters // 16, initializer=initializer)
          self.se_128 = SEBlock(filters // 16, initializer=initializer)
          
          self.up_256 = upBlock(filters // 32, initializer=initializer)
          self.se_256 = SEBlock(filters // 32, initializer=initializer)
          self.ch_conv = layers.Conv2D(3, 3, padding='same', 
                          use_bias=False, kernel_initializer=initializer)
          self.tanh = layers.Activation('tanh', dtype='float32')
                         
      def call(self, z):
          z = normalize_2nd_moment(z)
          feat_4 = self.init(z)
          feat_8 = self.up_8(feat_4)   
          feat_16 = self.up_16(feat_8) 
          feat_32 = self.up_32(feat_16) 
          feat_64 = self.up_64(feat_32) * self.se_64(feat_4)
          feat_128 = self.up_128(feat_64) * self.se_128(feat_8)
          feat_256 = self.up_256(feat_128) * self.se_256(feat_16)
          img_256 = self.ch_conv(feat_256)
          return [self.tanh(img_256)]
  
  
  class DownBlockComp(layers.Layer):
      def __init__(self, filters, initializer='orthogonal'):
          super(DownBlockComp, self).__init__()
  
          self.main = tf.keras.Sequential([
              layers.Conv2D(filters, kernel_size=4, padding='same',
                  strides=2, use_bias=False, kernel_initializer=initializer),
              layers.BatchNormalization(),
              layers.LeakyReLU(0.2),
              layers.Conv2D(filters, kernel_size=3, padding='same',
                  use_bias=False, kernel_initializer=initializer),
              layers.BatchNormalization(),
              layers.LeakyReLU(0.2),
          ])
          
          self.direct = tf.keras.Sequential([
              layers.AveragePooling2D((2, 2)),
              layers.Conv2D(filters, kernel_size=1, padding='same',
                  use_bias=False, kernel_initializer=initializer),
              layers.BatchNormalization(),
              layers.LeakyReLU(0.2),
          ])
  
      def call(self, x):
          return (self.main(x) + self.direct(x)) / 2
  
  
  def decode_image(filters=128, initializer='orthogonal'):
      encode_img = tf.keras.Sequential([
          AdaptiveAveragePooling2D(8),
          upBlock(filters, initializer=initializer),
          upBlock(filters // 2, initializer=initializer),
          upBlock(filters // 4, initializer=initializer),
          upBlock(filters // 8, initializer=initializer),
          layers.Conv2D(3, 3, padding='same', 
              use_bias=False, kernel_initializer=initializer),
          layers.Activation('tanh', dtype='float32')
      ])
      return encode_img
  
  
  class Discriminator(tf.keras.models.Model):
      def __init__(self, filters=128, initializer='orthogonal',
                   dec_dim=128):
          super(Discriminator, self).__init__()
          '''Encode image'''
          self.down_from_big = tf.keras.Sequential([
              layers.Conv2D(filters // 32, kernel_size=3, padding='same', 
                  use_bias=False, kernel_initializer=initializer),
              layers.LeakyReLU(0.2)
          ])
          self.down_128 = DownBlockComp(filters // 16, initializer=initializer)
          self.down_64 = DownBlockComp(filters // 8, initializer=initializer)
          self.down_32 = DownBlockComp(filters // 4, initializer=initializer)
          self.down_16 = DownBlockComp(filters // 2, initializer=initializer)
          self.down_8 = DownBlockComp(filters, initializer=initializer)
  
          '''Logits'''
          self.logits = tf.keras.Sequential([
              layers.Conv2D(filters, kernel_size=1, padding='valid', 
                  use_bias=False, kernel_initializer=initializer),  
              layers.LeakyReLU(0.2),
              layers.Conv2D(1, kernel_size=4, padding='valid', 
                  use_bias=False, kernel_initializer=initializer), 
              layers.Flatten()
          ])
          self.decoder = decode_image(dec_dim, initializer=initializer)
   
      def call(self, img, decode=False):
          feat_256 = self.down_from_big(img)
          feat_128 = self.down_128(feat_256)  
          feat_64 = self.down_64(feat_128)  
          feat_32 = self.down_32(feat_64)  
          feat_16 = self.down_16(feat_32)
          feat_8 = self.down_8(feat_16)
  
          if decode:
              return [self.logits(feat_8), self.decoder(feat_8)]
          else:
              return [self.logits(feat_8)]
  
  