                                  name="{}_seperableconv2d".format(name),
                                  **kwargs)(inp)
          var_x = Activation("relu", name="{}_relu".format(name))(var_x)
          return var_x
  
  # <<< GAN V2.2 Blocks >>> #
  # TODO Merge these into NNBLock class when porting GAN2.2
  
  
  # Gan Constansts:
  GAN22_CONV_INIT = "he_normal"
  GAN22_REGULARIZER = 1e-4
  
  
  # Gan Blocks:
  def normalization(inp, norm="none", group="16"):
      """ GAN Normalization """
      if norm == "layernorm":
          var_x = GroupNormalization(group=group)(inp)
      elif norm == "batchnorm":
          var_x = BatchNormalization()(inp)
      elif norm == "groupnorm":
          var_x = GroupNormalization(group=16)(inp)
      elif norm == "instancenorm":
          var_x = InstanceNormalization()(inp)
      elif norm == "hybrid":
          if group % 2 == 1:
              raise ValueError("Output channels must be an even number for hybrid norm, "
                               "received {}.".format(group))
          filt = group
          var_x_0 = Lambda(lambda var_x: var_x[..., :filt // 2])(var_x)
          var_x_1 = Lambda(lambda var_x: var_x[..., filt // 2:])(var_x)
          var_x_0 = Conv2D(filt // 2,
                           kernel_size=1,
                           kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                           kernel_initializer=GAN22_CONV_INIT)(var_x_0)
          var_x_1 = InstanceNormalization()(var_x_1)
          var_x = concatenate([var_x_0, var_x_1], axis=-1)
      else:
          var_x = inp
      return var_x
  
  
  def upscale_ps(inp, filters, initializer, use_norm=False, norm="none"):
      """ GAN Upscaler - Pixel Shuffler """
      var_x = Conv2D(filters * 4,
                     kernel_size=3,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                     kernel_initializer=initializer,
                     padding="same")(inp)
      var_x = LeakyReLU(0.2)(var_x)
      var_x = normalization(var_x, norm, filters) if use_norm else var_x
      var_x = PixelShuffler()(var_x)
      return var_x
  
  
  def upscale_nn(inp, filters, use_norm=False, norm="none"):
      """ GAN Neural Network """
      var_x = UpSampling2D()(inp)
      var_x = reflect_padding_2d(var_x, 1)
      var_x = Conv2D(filters,
                     kernel_size=3,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                     kernel_initializer="he_normal")(var_x)
      var_x = normalization(var_x, norm, filters) if use_norm else var_x
      return var_x
  
  
  def reflect_padding_2d(inp, pad=1):
      """ GAN Reflect Padding (2D) """
      var_x = Lambda(lambda var_x: tf.pad(var_x,
                                          [[0, 0], [pad, pad], [pad, pad], [0, 0]],
                                          mode="REFLECT"))(inp)
      return var_x
  
  
  def conv_gan(inp, filters, use_norm=False, strides=2, norm="none"):
      """ GAN Conv Block """
      var_x = Conv2D(filters,
                     kernel_size=3,
                     strides=strides,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                     kernel_initializer=GAN22_CONV_INIT,
                     use_bias=False,
                     padding="same")(inp)
      var_x = Activation("relu")(var_x)
      var_x = normalization(var_x, norm, filters) if use_norm else var_x
      return var_x
  
  
  def conv_d_gan(inp, filters, use_norm=False, norm="none"):
      """ GAN Discriminator Conv Block """
      var_x = inp
      var_x = Conv2D(filters,
                     kernel_size=4,
                     strides=2,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                     kernel_initializer=GAN22_CONV_INIT,
                     use_bias=False,
                     padding="same")(var_x)
      var_x = LeakyReLU(alpha=0.2)(var_x)
      var_x = normalization(var_x, norm, filters) if use_norm else var_x
      return var_x
  
  
  def res_block_gan(inp, filters, use_norm=False, norm="none"):
      """ GAN Res Block """
      var_x = Conv2D(filters,
                     kernel_size=3,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                     kernel_initializer=GAN22_CONV_INIT,
                     use_bias=False,
                     padding="same")(inp)
      var_x = LeakyReLU(alpha=0.2)(var_x)
      var_x = normalization(var_x, norm, filters) if use_norm else var_x
      var_x = Conv2D(filters,
                     kernel_size=3,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER),
                     kernel_initializer=GAN22_CONV_INIT,
                     use_bias=False,
                     padding="same")(var_x)
      var_x = add([var_x, inp])
      var_x = LeakyReLU(alpha=0.2)(var_x)
      var_x = normalization(var_x, norm, filters) if use_norm else var_x
      return var_x
  
  
  def self_attn_block(inp, n_c, squeeze_factor=8):
      """ GAN Self Attention Block
      Code borrows from https://github.com/taki0112/Self-Attention-GAN-Tensorflow
      """
      msg = "Input channels must be >= {}, recieved nc={}".format(squeeze_factor, n_c)
      assert n_c // squeeze_factor > 0, msg
      var_x = inp
      shape_x = var_x.get_shape().as_list()
  
      var_f = Conv2D(n_c // squeeze_factor, 1,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER))(var_x)
      var_g = Conv2D(n_c // squeeze_factor, 1,
                     kernel_regularizer=regularizers.l2(GAN22_REGULARIZER))(var_x)
      var_h = Conv2D(n_c, 1, kernel_regularizer=regularizers.l2(GAN22_REGULARIZER))(var_x)
  
      shape_f = var_f.get_shape().as_list()
      shape_g = var_g.get_shape().as_list()
      shape_h = var_h.get_shape().as_list()
      flat_f = Reshape((-1, shape_f[-1]))(var_f)
      flat_g = Reshape((-1, shape_g[-1]))(var_g)
      flat_h = Reshape((-1, shape_h[-1]))(var_h)
  
      var_s = Lambda(lambda var_x: K.batch_dot(var_x[0],
                                               Permute((2, 1))(var_x[1])))([flat_g, flat_f])
  
      beta = Softmax(axis=-1)(var_s)
      var_o = Lambda(lambda var_x: K.batch_dot(var_x[0], var_x[1]))([beta, flat_h])
      var_o = Reshape(shape_x[1:])(var_o)
      var_o = Scale()(var_o)
  
      out = add([var_o, inp])
      return out
 
</s> remove                                 **kwargs)(inp)
 </s> add                                 **kwargs)(input_tensor) </s> remove             self.switch_kernel_initializer(kwargs, original_init)
        var_x = Add()([var_x, inp])
 </s> add             self._switch_kernel_initializer(kwargs, original_init)
        var_x = Add()([var_x, input_tensor]) </s> remove class GroupNormalization(Layer):
    """ Group Normalization
        from: shoanlu GAN: https://github.com/shaoanlu/faceswap-GAN"""

    def __init__(self, axis=-1,
                 gamma_init='one', beta_init='zero',
                 gamma_regularizer=None, beta_regularizer=None,
                 epsilon=1e-6,
                 group=32,
                 data_format=None,
                 **kwargs):
        self.beta = None
        self.gamma = None
        super(GroupNormalization, self).__init__(**kwargs)

        self.axis = to_list(axis)
        self.gamma_init = initializers.get(gamma_init)
        self.beta_init = initializers.get(beta_init)
        self.gamma_regularizer = regularizers.get(gamma_regularizer)
        self.beta_regularizer = regularizers.get(beta_regularizer)
        self.epsilon = epsilon
        self.group = group
        self.data_format = K.normalize_data_format(data_format)

        self.supports_masking = True

    def build(self, input_shape):
        self.input_spec = [InputSpec(shape=input_shape)]
        shape = [1 for _ in input_shape]
        if self.data_format == 'channels_last':
            channel_axis = -1
            shape[channel_axis] = input_shape[channel_axis]
        elif self.data_format == 'channels_first':
            channel_axis = 1
            shape[channel_axis] = input_shape[channel_axis]
        # for i in self.axis:
        #    shape[i] = input_shape[i]
        self.gamma = self.add_weight(shape=shape,
                                     initializer=self.gamma_init,
                                     regularizer=self.gamma_regularizer,
                                     name='gamma')
        self.beta = self.add_weight(shape=shape,
                                    initializer=self.beta_init,
                                    regularizer=self.beta_regularizer,
                                    name='beta')
        self.built = True

    def call(self, inputs, mask=None):
        input_shape = K.int_shape(inputs)
        if len(input_shape) != 4 and len(input_shape) != 2:
            raise ValueError('Inputs should have rank ' +
                             str(4) + " or " + str(2) +
                             '; Received input shape:', str(input_shape))

        if len(input_shape) == 4:
            if self.data_format == 'channels_last':
                batch_size, height, width, channels = input_shape
                if batch_size is None:
                    batch_size = -1

                if channels < self.group:
                    raise ValueError('Input channels should be larger than group size' +
                                     '; Received input channels: ' + str(channels) +
                                     '; Group size: ' + str(self.group))

                var_x = K.reshape(inputs, (batch_size,
                                           height,
                                           width,
                                           self.group,
                                           channels // self.group))
                mean = K.mean(var_x, axis=[1, 2, 4], keepdims=True)
                std = K.sqrt(K.var(var_x, axis=[1, 2, 4], keepdims=True) + self.epsilon)
                var_x = (var_x - mean) / std

                var_x = K.reshape(var_x, (batch_size, height, width, channels))
                retval = self.gamma * var_x + self.beta
            elif self.data_format == 'channels_first':
                batch_size, channels, height, width = input_shape
                if batch_size is None:
                    batch_size = -1

                if channels < self.group:
                    raise ValueError('Input channels should be larger than group size' +
                                     '; Received input channels: ' + str(channels) +
                                     '; Group size: ' + str(self.group))

                var_x = K.reshape(inputs, (batch_size,
                                           self.group,
                                           channels // self.group,
                                           height,
                                           width))
                mean = K.mean(var_x, axis=[2, 3, 4], keepdims=True)
                std = K.sqrt(K.var(var_x, axis=[2, 3, 4], keepdims=True) + self.epsilon)
                var_x = (var_x - mean) / std

                var_x = K.reshape(var_x, (batch_size, channels, height, width))
                retval = self.gamma * var_x + self.beta

        elif len(input_shape) == 2:
            reduction_axes = list(range(0, len(input_shape)))
            del reduction_axes[0]
            batch_size, _ = input_shape
            if batch_size is None:
                batch_size = -1

            mean = K.mean(inputs, keepdims=True)
            std = K.sqrt(K.var(inputs, keepdims=True) + self.epsilon)
            var_x = (inputs - mean) / std

            retval = self.gamma * var_x + self.beta
        return retval

    def get_config(self):
        config = {'epsilon': self.epsilon,
                  'axis': self.axis,
                  'gamma_init': initializers.serialize(self.gamma_init),
                  'beta_init': initializers.serialize(self.beta_init),
                  'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
                  'beta_regularizer': regularizers.serialize(self.gamma_regularizer),
                  'group': self.group}
        base_config = super(GroupNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


# Update normalizations into Keras custom objects
 </s> add # Update normalization into Keras custom objects </s> remove             kwargs["name"] = self.get_name("conv2d_{}".format(inp.shape[1]))
        kwargs = self.set_default_initializer(kwargs)
 </s> add             kwargs["name"] = self._get_name("conv2d_{}".format(input_tensor.shape[1]))
        kwargs = self._set_default_initializer(kwargs) </s> remove                        **kwargs)(inp)
 </s> add                        **kwargs)(input_tensor)
                                  name="{}_seperableconv2d".format(name),
                                  **kwargs)(inp)
          var_x = Activation("relu", name="{}_relu".format(name))(var_x)
          return var_x



























































































































































keep keep keep keep replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace replace
 Simple backend unit tests (#1020)

* Add simple backend tests for lib.model
* Document lib.model
* Fix GMSD Loss for AMD
* Remove obsolete code from lib.model 
lib/model/nn_blocks.py