networks/layers.py

import tensorflow as tf
import numpy as np
import scipy.signal as sps
import scipy.special as spspec

import tensorflow.keras.backend as K
import math
from tensorflow.python.keras.utils import conv_utils
from tensorflow.keras.layers import Layer, InputSpec, Conv2D, LeakyReLU, Dense, BatchNormalization, Input, Concatenate
from tensorflow.keras.layers import Conv2DTranspose, ReLU, Activation, UpSampling2D, Add, Reshape, Multiply
from tensorflow.keras.layers import AveragePooling2D, LayerNormalization, GlobalAveragePooling2D, MaxPooling2D, Flatten
from tensorflow.keras import initializers, constraints, regularizers
from tensorflow.keras.models import Model
from tensorflow_addons.layers import InstanceNormalization


def sin_activation(x, omega=30):
 return tf.math.sin(omega * x)


class AdaIN(Layer):
 def __init__(self, **kwargs):
 super(AdaIN, self).__init__(**kwargs)

 def build(self, input_shapes):
 x_shape = input_shapes[0]
 w_shape = input_shapes[1]

 self.w_channels = w_shape[-1]
 self.x_channels = x_shape[-1]

 self.dense_1 = Dense(self.x_channels)
 self.dense_2 = Dense(self.x_channels)

 def call(self, inputs):
 x, w = inputs
 ys = tf.reshape(self.dense_1(w), (-1, 1, 1, self.x_channels))
 yb = tf.reshape(self.dense_2(w), (-1, 1, 1, self.x_channels))
 return ys * x + yb

 def get_config(self):
 config = {
 #'w_channels': self.w_channels,
 #'x_channels': self.x_channels
 }
 base_config = super(AdaIN, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class Conv2DMod(Layer):

 def __init__(self,
 filters,
 kernel_size,
 strides=1,
 padding='valid',
 dilation_rate=1,
 kernel_initializer='glorot_uniform',
 kernel_regularizer=None,
 activity_regularizer=None,
 kernel_constraint=None,
 demod=True,
 **kwargs):
 super(Conv2DMod, self).__init__(**kwargs)
 self.filters = filters
 self.rank = 2
 self.kernel_size = conv_utils.normalize_tuple(kernel_size, 2, 'kernel_size')
 self.strides = conv_utils.normalize_tuple(strides, 2, 'strides')
 self.padding = conv_utils.normalize_padding(padding)
 self.dilation_rate = conv_utils.normalize_tuple(dilation_rate, 2, 'dilation_rate')
 self.kernel_initializer = initializers.get(kernel_initializer)
 self.kernel_regularizer = regularizers.get(kernel_regularizer)
 self.activity_regularizer = regularizers.get(activity_regularizer)
 self.kernel_constraint = constraints.get(kernel_constraint)
 self.demod = demod
 self.input_spec = [InputSpec(ndim = 4),
 InputSpec(ndim = 2)]

 def build(self, input_shape):
 channel_axis = -1
 if input_shape[0][channel_axis] is None:
 raise ValueError('The channel dimension of the inputs '
 'should be defined. Found `None`.')
 input_dim = input_shape[0][channel_axis]
 kernel_shape = self.kernel_size + (input_dim, self.filters)

 if input_shape[1][-1] != input_dim:
 raise ValueError('The last dimension of modulation input should be equal to input dimension.')

 self.kernel = self.add_weight(shape=kernel_shape,
 initializer=self.kernel_initializer,
 name='kernel',
 regularizer=self.kernel_regularizer,
 constraint=self.kernel_constraint)

 # Set input spec.
 self.input_spec = [InputSpec(ndim=4, axes={channel_axis: input_dim}),
 InputSpec(ndim=2)]
 self.built = True

 def call(self, inputs):

 #To channels last
 x = tf.transpose(inputs[0], [0, 3, 1, 2])

 #Get weight and bias modulations
 #Make sure w's shape is compatible with self.kernel
 w = K.expand_dims(K.expand_dims(K.expand_dims(inputs[1], axis = 1), axis = 1), axis = -1)

 #Add minibatch layer to weights
 wo = K.expand_dims(self.kernel, axis = 0)

 #Modulate
 weights = wo * (w+1)

 #Demodulate
 if self.demod:
 d = K.sqrt(K.sum(K.square(weights), axis=[1,2,3], keepdims = True) + 1e-8)
 weights = weights / d

 #Reshape/scale input
 x = tf.reshape(x, [1, -1, x.shape[2], x.shape[3]]) # Fused => reshape minibatch to convolution groups.
 w = tf.reshape(tf.transpose(weights, [1, 2, 3, 0, 4]), [weights.shape[1], weights.shape[2], weights.shape[3], -1])

 x = tf.nn.conv2d(x, w,
 strides=self.strides,
 padding="SAME",
 data_format="NCHW")

 # Reshape/scale output.
 x = tf.reshape(x, [-1, self.filters, x.shape[2], x.shape[3]]) # Fused => reshape convolution groups back to minibatch.
 x = tf.transpose(x, [0, 2, 3, 1])

 return x

 def compute_output_shape(self, input_shape):
 space = input_shape[0][1:-1]
 new_space = []
 for i in range(len(space)):
 new_dim = conv_utils.conv_output_length(
 space[i],
 self.kernel_size[i],
 padding=self.padding,
 stride=self.strides[i],
 dilation=self.dilation_rate[i])
 new_space.append(new_dim)

 return (input_shape[0],) + tuple(new_space) + (self.filters,)

 def get_config(self):
 config = {
 'filters': self.filters,
 'kernel_size': self.kernel_size,
 'strides': self.strides,
 'padding': self.padding,
 'dilation_rate': self.dilation_rate,
 'kernel_initializer': initializers.serialize(self.kernel_initializer),
 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer),
 'activity_regularizer':
 regularizers.serialize(self.activity_regularizer),
 'kernel_constraint': constraints.serialize(self.kernel_constraint),
 'demod': self.demod
 }
 base_config = super(Conv2DMod, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class CreatePatches(tf.keras.layers.Layer ):

 def __init__( self , patch_size):
 super( CreatePatches , self).__init__()
 self.patch_size = patch_size

 def call(self, inputs):
 patches = []
 # For square images only ( as inputs.shape[ 1 ] = inputs.shape[ 2 ] )
 input_image_size = inputs.shape[ 1 ]
 for i in range( 0 , input_image_size , self.patch_size ):
 for j in range( 0 , input_image_size , self.patch_size ):
 patches.append( inputs[ : , i : i + self.patch_size , j : j + self.patch_size , : ] )
 return patches

 def get_config(self):
 config = {'patch_size': self.patch_size,
 }
 base_config = super(CreatePatches, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class SelfAttention(tf.keras.layers.Layer ):

 def __init__( self , alpha, filters=128):
 super(SelfAttention , self).__init__()
 self.alpha = alpha
 self.filters = filters

 self.f = Conv2D(filters, 1, 1)
 self.g = Conv2D(filters, 1, 1)
 self.s = Conv2D(filters, 1, 1)

 def call(self, inputs):

 f_map = self.f(inputs)
 f_map = tf.image.transpose(f_map)

 g_map = self.g(inputs)

 s_map = self.s(inputs)

 att = f_map * g_map

 att = att / self.alpha

 return tf.keras.activations.softmax(att + s_map, axis=0)

 def get_config(self):
 config = {'alpha': self.alpha,
 }
 base_config = super(SelfAttention, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class Sampling(Layer):
 """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

 def call(self, inputs):
 z_mean, z_log_var = inputs
 batch = tf.shape(z_mean)[0]
 dim = tf.shape(z_mean)[1]
 epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
 return z_mean + tf.exp(0.5 * z_log_var) * epsilon

 def get_config(self):
 config = {
 }
 base_config = super(Sampling, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class ArcMarginPenaltyLogists(tf.keras.layers.Layer):
 """ArcMarginPenaltyLogists"""
 def __init__(self, num_classes, margin=0.7, logist_scale=64, **kwargs):
 super(ArcMarginPenaltyLogists, self).__init__(**kwargs)
 self.num_classes = num_classes
 self.margin = margin
 self.logist_scale = logist_scale

 def build(self, input_shape):
 self.w = self.add_variable(
 "weights", shape=[int(input_shape[-1]), self.num_classes])
 self.cos_m = tf.identity(math.cos(self.margin), name='cos_m')
 self.sin_m = tf.identity(math.sin(self.margin), name='sin_m')
 self.th = tf.identity(math.cos(math.pi - self.margin), name='th')
 self.mm = tf.multiply(self.sin_m, self.margin, name='mm')

 def call(self, embds, labels):
 normed_embds = tf.nn.l2_normalize(embds, axis=1, name='normed_embd')
 normed_w = tf.nn.l2_normalize(self.w, axis=0, name='normed_weights')

 cos_t = tf.matmul(normed_embds, normed_w, name='cos_t')
 sin_t = tf.sqrt(1. - cos_t ** 2, name='sin_t')

 cos_mt = tf.subtract(
 cos_t * self.cos_m, sin_t * self.sin_m, name='cos_mt')

 cos_mt = tf.where(cos_t > self.th, cos_mt, cos_t - self.mm)

 mask = tf.one_hot(tf.cast(labels, tf.int32), depth=self.num_classes,
 name='one_hot_mask')

 logists = tf.where(mask == 1., cos_mt, cos_t)
 logists = tf.multiply(logists, self.logist_scale, 'arcface_logist')

 return logists

 def get_config(self):
 config = {'num_classes': self.num_classes,
 'margin': self.margin,
 'logist_scale': self.logist_scale
 }
 base_config = super(ArcMarginPenaltyLogists, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))

class KLLossLayer(tf.keras.layers.Layer):
 """ArcMarginPenaltyLogists"""
 def __init__(self, beta=1.5, **kwargs):
 super(KLLossLayer, self).__init__(**kwargs)
 self.beta = beta

 def call(self, inputs):
 z_mean, z_log_var = inputs

 kl_loss = tf.reduce_mean(1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
 kl_loss = -0.5 * kl_loss * self.beta

 self.add_loss(kl_loss * 0)
 self.add_metric(kl_loss, 'kl_loss')

 return inputs

 def get_config(self):
 config = {
 'beta': self.beta
 }
 base_config = super(KLLossLayer, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class ReflectionPadding2D(Layer):
 def __init__(self, padding=(1, 1), **kwargs):
 self.padding = tuple(padding)
 self.input_spec = [InputSpec(ndim=4)]
 super(ReflectionPadding2D, self).__init__(**kwargs)

 def compute_output_shape(self, s):
 return (s[0], s[1] + 2 * self.padding[0], s[2] + 2 * self.padding[1], s[3])

 def call(self, x, mask=None):
 w_pad,h_pad = self.padding
 return tf.pad(x, [[0,0], [h_pad,h_pad], [w_pad,w_pad], [0,0] ], 'REFLECT')

 def get_config(self):
 config = {
 'padding': self.padding,
 }
 base_config = super(ReflectionPadding2D, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class ResBlock(Layer):

 def __init__(self, fil, **kwargs):
 super(ResBlock, self).__init__(**kwargs)
 self.fil = fil

 self.conv_0 = Conv2D(kernel_size=3, filters=fil, strides=1)
 self.conv_1 = Conv2D(kernel_size=3, filters=fil, strides=1)

 self.res = Conv2D(kernel_size=1, filters=1, strides=1)

 self.lrelu = LeakyReLU(0.2)
 self.padding = ReflectionPadding2D(padding=(1, 1))

 def call(self, inputs):
 res = self.res(inputs)

 x = self.padding(inputs)
 x = self.conv_0(x)
 x = self.lrelu(x)

 x = self.padding(x)
 x = self.conv_1(x)
 x = self.lrelu(x)

 out = x + res

 return out

 def get_config(self):
 config = {
 'fil': self.fil
 }
 base_config = super(ResBlock, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))


class SubpixelConv2D(Layer):
 """ Subpixel Conv2D Layer
 upsampling a layer from (h, w, c) to (h*r, w*r, c/(r*r)),
 where r is the scaling factor, default to 4
 # Arguments
 upsampling_factor: the scaling factor
 # Input shape
 Arbitrary. Use the keyword argument `input_shape`
 (tuple of integers, does not include the samples axis)
 when using this layer as the first layer in a model.
 # Output shape
 the second and the third dimension increased by a factor of
 `upsampling_factor`; the last layer decreased by a factor of
 `upsampling_factor^2`.
 # References
 Real-Time Single Image and Video Super-Resolution Using an Efficient
 Sub-Pixel Convolutional Neural Network Shi et Al. https://arxiv.org/abs/1609.05158
 """

 def __init__(self, upsampling_factor=4, **kwargs):
 super(SubpixelConv2D, self).__init__(**kwargs)
 self.upsampling_factor = upsampling_factor

 def build(self, input_shape):
 last_dim = input_shape[-1]
 factor = self.upsampling_factor * self.upsampling_factor
 if last_dim % (factor) != 0:
 raise ValueError('Channel ' + str(last_dim) + ' should be of '
 'integer times of upsampling_factor^2: ' +
 str(factor) + '.')

 def call(self, inputs, **kwargs):
 return tf.nn.depth_to_space( inputs, self.upsampling_factor )

 def get_config(self):
 config = { 'upsampling_factor': self.upsampling_factor, }
 base_config = super(SubpixelConv2D, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))

 def compute_output_shape(self, input_shape):
 factor = self.upsampling_factor * self.upsampling_factor
 input_shape_1 = None
 if input_shape[1] is not None:
 input_shape_1 = input_shape[1] * self.upsampling_factor
 input_shape_2 = None
 if input_shape[2] is not None:
 input_shape_2 = input_shape[2] * self.upsampling_factor
 dims = [ input_shape[0],
 input_shape_1,
 input_shape_2,
 int(input_shape[3]/factor)
 ]
 return tuple( dims )


def id_mod_res(inputs, c):
 feature_map, z_id = inputs

 x = Conv2D(c, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.0001))(feature_map)

 x = AdaIN()([x, z_id])

 x = ReLU()(x)

 x = Conv2D(c, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.0001))(x)

 x = AdaIN()([x, z_id])

 out = Add()([x, feature_map])

 return out


def id_mod_res_v2(inputs, c):
 feature_map, z_id = inputs

 affine = Dense(feature_map.shape[-1])(z_id)
 x = Conv2DMod(c, kernel_size=3, padding='same',
 kernel_regularizer=tf.keras.regularizers.l2(0.0001))([feature_map, affine])

 x = ReLU()(x)

 affine = Dense(x.shape[-1])(z_id)
 x = Conv2DMod(c, kernel_size=3, padding='same',
 kernel_regularizer=tf.keras.regularizers.l2(0.0001))([x, affine])
 out = Add()([x, feature_map])

 x = ReLU()(x)

 return out


def simswap(im_size, filter_scale=1, deep=True):
 inputs = Input(shape=(im_size, im_size, 3))
 z_id = Input(shape=(512,))

 x = ReflectionPadding2D(padding=(3, 3))(inputs)
 x = Conv2D(filters=64 // filter_scale, kernel_size=7, padding='valid', kernel_regularizer=tf.keras.regularizers.l1(0.0001))(x) # 112
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = Conv2D(filters=64 // filter_scale, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.0001))(x) # 56
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = Conv2D(filters=256 // filter_scale, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.0001))(x) # 28
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = Conv2D(filters=512 // filter_scale, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.0001))(x) # 14
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x) # 14

 if deep:
 x = Conv2D(filters=512 // filter_scale, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.0001))(x) # 7
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = id_mod_res([x, z_id], 512 // filter_scale)

 x = id_mod_res([x, z_id], 512 // filter_scale)

 x = id_mod_res([x, z_id], 512 // filter_scale)

 x = id_mod_res([x, z_id], 512 // filter_scale)

 if deep:
 x = SubpixelConv2D(upsampling_factor=2)(x)
 x = Conv2D(filters=512 // filter_scale, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = SubpixelConv2D(upsampling_factor=2)(x)
 x = Conv2D(filters=256 // filter_scale, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = SubpixelConv2D(upsampling_factor=2)(x)
 x = Conv2D(filters=128 // filter_scale, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x) # 56

 x = SubpixelConv2D(upsampling_factor=2)(x)
 x = Conv2D(filters=64 // filter_scale, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l1(0.001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x) # 112

 x = ReflectionPadding2D(padding=(3, 3))(x)
 out = Conv2D(filters=3, kernel_size=7, padding='valid')(x) # 112

 model = Model([inputs, z_id], out)
 model.summary()

 return model


def simswap_v2(deep=True):
 inputs = Input(shape=(224, 224, 3))
 z_id = Input(shape=(512,))

 x = ReflectionPadding2D(padding=(3, 3))(inputs)
 x = Conv2D(filters=64, kernel_size=7, padding='valid', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x) # 112
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = Conv2D(filters=64, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x) # 56
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = Conv2D(filters=256, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x) # 28
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = Conv2D(filters=512, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x) # 14
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x) # 14

 if deep:
 x = Conv2D(filters=512, kernel_size=3, strides=2, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x) # 7
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = id_mod_res_v2([x, z_id], 512)

 x = id_mod_res_v2([x, z_id], 512)

 x = id_mod_res_v2([x, z_id], 512)

 x = id_mod_res_v2([x, z_id], 512)

 x = id_mod_res_v2([x, z_id], 512)

 x = id_mod_res_v2([x, z_id], 512)

 if deep:
 x = UpSampling2D(interpolation='bilinear')(x)
 x = Conv2D(filters=512, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = UpSampling2D(interpolation='bilinear')(x)
 x = Conv2D(filters=256, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x)

 x = UpSampling2D(interpolation='bilinear')(x)
 x = Conv2D(filters=128, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x) # 56

 x = UpSampling2D(interpolation='bilinear')(x)
 x = Conv2D(filters=64, kernel_size=3, padding='same', kernel_regularizer=tf.keras.regularizers.l2(0.0001))(x)
 x = BatchNormalization()(x)
 x = Activation(tf.keras.activations.relu)(x) # 112

 x = ReflectionPadding2D(padding=(3, 3))(x)
 out = Conv2D(filters=3, kernel_size=7, padding='valid')(x) # 112
 out = Activation('sigmoid')(out)

 model = Model([inputs, z_id], out)
 model.summary()

 return model


class AdaptiveAttention(Layer):

 def __init__(self, **kwargs):
 super(AdaptiveAttention, self).__init__(**kwargs)

 def call(self, inputs):
 m, a, i = inputs
 return (1 - m) * a + m * i

 def get_config(self):
 base_config = super(AdaptiveAttention, self).get_config()
 return base_config


def aad_block(inputs, c_out):
 h, z_att, z_id = inputs

 h_norm = BatchNormalization()(h)
 h = Conv2D(filters=c_out, kernel_size=1, kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(h_norm)

 m = Activation('sigmoid')(h)

 z_att_gamma = Conv2D(filters=c_out,
 kernel_size=1,
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(z_att)

 z_att_beta = Conv2D(filters=c_out,
 kernel_size=1,
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(z_att)

 a = Multiply()([h_norm, z_att_gamma])
 a = Add()([a, z_att_beta])

 z_id_gamma = Dense(h_norm.shape[-1],
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(z_id)
 z_id_gamma = Reshape(target_shape=(1, 1, h_norm.shape[-1]))(z_id_gamma)

 z_id_beta = Dense(h_norm.shape[-1],
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(z_id)
 z_id_beta = Reshape(target_shape=(1, 1, h_norm.shape[-1]))(z_id_beta)

 i = Multiply()([h_norm, z_id_gamma])
 i = Add()([i, z_id_beta])

 h_out = AdaptiveAttention()([m, a, i])

 return h_out


def aad_block_mod(inputs, c_out):
 h, z_att, z_id = inputs

 h_norm = BatchNormalization()(h)
 h = Conv2D(filters=c_out, kernel_size=1, kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(h_norm)

 m = Activation('sigmoid')(h)

 z_att_gamma = Conv2D(filters=c_out,
 kernel_size=1,
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(z_att)

 z_att_beta = Conv2D(filters=c_out,
 kernel_size=1,
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(z_att)

 a = Multiply()([h_norm, z_att_gamma])
 a = Add()([a, z_att_beta])

 z_id_gamma = Dense(h_norm.shape[-1],
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(z_id)

 i = Conv2DMod(filters=c_out,
 kernel_size=1,
 padding='same',
 kernel_initializer='he_uniform',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))([h_norm, z_id_gamma])

 h_out = AdaptiveAttention()([m, a, i])

 return h_out


def aad_res_block(inputs, c_in, c_out):
 h, z_att, z_id = inputs

 if c_in == c_out:
 aad = aad_block([h, z_att, z_id], c_out)
 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(act)

 aad = aad_block([conv, z_att, z_id], c_out)
 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(act)

 h_out = Add()([h, conv])
 return h_out
 else:
 aad = aad_block([h, z_att, z_id], c_in)
 act = ReLU()(aad)
 h_res = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(act)

 aad = aad_block([h, z_att, z_id], c_in)
 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(act)

 aad = aad_block([conv, z_att, z_id], c_out)
 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.001))(act)

 h_out = Add()([h_res, conv])

 return h_out


def aad_res_block_mod(inputs, c_in, c_out):
 h, z_att, z_id = inputs

 if c_in == c_out:
 aad = aad_block_mod([h, z_att, z_id], c_out)

 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(act)

 aad = aad_block_mod([conv, z_att, z_id], c_out)
 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(act)

 h_out = Add()([h, conv])
 return h_out
 else:
 aad = aad_block_mod([h, z_att, z_id], c_in)
 act = ReLU()(aad)
 h_res = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(act)

 aad = aad_block_mod([h, z_att, z_id], c_in)
 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(act)

 aad = aad_block_mod([conv, z_att, z_id], c_out)
 act = ReLU()(aad)
 conv = Conv2D(filters=c_out,
 kernel_size=3,
 padding='same',
 kernel_regularizer=tf.keras.regularizers.l1(l1=0.0001))(act)

 h_out = Add()([h_res, conv])

 return h_out


class FilteredReLU(Layer):

 def __init__(self,
 critically_sampled,

 in_channels,
 out_channels,
 in_size,
 out_size,
 in_sampling_rate,
 out_sampling_rate,
 in_cutoff,
 out_cutoff,
 in_half_width,
 out_half_width,

 conv_kernel = 3,
 lrelu_upsampling = 2,
 filter_size = 6,
 conv_clamp = 256,
 use_radial_filters = False,
 is_torgb = False,
 **kwargs):
 super(FilteredReLU, self).__init__(**kwargs)
 self.critically_sampled = critically_sampled

 self.in_channels = in_channels
 self.out_channels = out_channels
 self.in_size = np.broadcast_to(np.asarray(in_size), [2])
 self.out_size = np.broadcast_to(np.asarray(out_size), [2])
 self.in_sampling_rate = in_sampling_rate
 self.out_sampling_rate = out_sampling_rate
 self.in_cutoff = in_cutoff
 self.out_cutoff = out_cutoff
 self.in_half_width = in_half_width
 self.out_half_width = out_half_width

 self.is_torgb = is_torgb

 self.conv_kernel = 1 if is_torgb else conv_kernel
 self.lrelu_upsampling = lrelu_upsampling
 self.conv_clamp = conv_clamp

 self.tmp_sampling_rate = max(in_sampling_rate, out_sampling_rate) * (1 if is_torgb else lrelu_upsampling)

 # Up sampling filter
 self.u_factor = int(np.rint(self.tmp_sampling_rate / self.in_sampling_rate))
 assert self.in_sampling_rate * self.u_factor == self.tmp_sampling_rate
 self.u_taps = filter_size * self.u_factor if self.u_factor > 1 and not self.is_torgb else 1
 self.u_filter = self.design_lowpass_filter(numtaps=self.u_taps,
 cutoff=self.in_cutoff,
 width=self.in_half_width*2,
 fs=self.tmp_sampling_rate)

 # Down sampling filter
 self.d_factor = int(np.rint(self.tmp_sampling_rate / self.out_sampling_rate))
 assert self.out_sampling_rate * self.d_factor == self.tmp_sampling_rate
 self.d_taps = filter_size * self.d_factor if self.d_factor > 1 and not self.is_torgb else 1
 self.d_radial = use_radial_filters and not self.critically_sampled
 self.d_filter = self.design_lowpass_filter(numtaps=self.d_taps,
 cutoff=self.out_cutoff,
 width=self.out_half_width*2,
 fs=self.tmp_sampling_rate,
 radial=self.d_radial)
 # Compute padding
 pad_total = (self.out_size - 1) * self.d_factor + 1
 pad_total -= (self.in_size + self.conv_kernel - 1) * self.u_factor
 pad_total += self.u_taps + self.d_taps - 2
 pad_lo = (pad_total + self.u_factor) // 2
 pad_hi = pad_total - pad_lo
 self.padding = [int(pad_lo[0]), int(pad_hi[0]), int(pad_lo[1]), int(pad_hi[1])]

 self.gain = 1 if self.is_torgb else np.sqrt(2)
 self.slope = 1 if self.is_torgb else 0.2

 self.act_funcs = {'linear':
 {'func': lambda x, **_: x,
 'def_alpha': 0,
 'def_gain': 1},
 'lrelu':
 {'func': lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha),
 'def_alpha': 0.2,
 'def_gain': np.sqrt(2)},
 }

 b_init = tf.zeros_initializer()
 self.bias = tf.Variable(initial_value=b_init(shape=(out_channels,),
 dtype="float32"),
 trainable=True)

 def design_lowpass_filter(self, numtaps, cutoff, width, fs, radial=False):
 if numtaps == 1:
 return None

 if not radial:
 f = sps.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs)
 return f

 x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
 r = np.hypot(*np.meshgrid(x, x))
 f = spspec.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
 beta = sps.kaiser_beta(sps.kaiser_atten(numtaps, width / (fs / 2)))
 w = np.kaiser(numtaps, beta)
 f *= np.outer(w, w)
 f /= np.sum(f)
 return f

 def get_filter_size(self, f):
 if f is None:
 return 1, 1
 assert 1 <= f.ndim <= 2
 return f.shape[-1], f.shape[0] # width, height

 def parse_padding(self, padding):
 if isinstance(padding, int):
 padding = [padding, padding]
 assert isinstance(padding, (list, tuple))
 assert all(isinstance(x, (int, np.integer)) for x in padding)
 padding = [int(x) for x in padding]
 if len(padding) == 2:
 px, py = padding
 padding = [px, px, py, py]
 px0, px1, py0, py1 = padding
 return px0, px1, py0, py1

 def bias_act(self, x, b=None, dim=3, act='linear', alpha=None, gain=None, clamp=None):
 spec = self.act_funcs[act]
 alpha = float(alpha if alpha is not None else spec['def_alpha'])
 gain = float(gain if gain is not None else spec['def_gain'])
 clamp = float(clamp if clamp is not None else -1)

 if b is not None:
 x = x + tf.reshape(b, shape=[-1 if i == dim else 1 for i in range(len(x.shape))])
 x = spec['func'](x, alpha=alpha)

 if gain != 1:
 x = x * gain

 if clamp >= 0:
 x = tf.clip_by_value(x, -clamp, clamp)
 return x

 def parse_scaling(self, scaling):
 if isinstance(scaling, int):
 scaling = [scaling, scaling]
 sx, sy = scaling
 assert sx >= 1 and sy >= 1
 return sx, sy

 def upfirdn2d(self, x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
 if f is None:
 f = tf.ones([1, 1], dtype=tf.float32)

 batch_size, in_height, in_width, num_channels = x.shape

 upx, upy = self.parse_scaling(up)
 downx, downy = self.parse_scaling(down)
 padx0, padx1, pady0, pady1 = self.parse_padding(padding)

 upW = in_width * upx + padx0 + padx1
 upH = in_height * upy + pady0 + pady1
 assert upW >= f.shape[-1] and upH >= f.shape[0]

 # Channel first format.
 x = tf.transpose(x, perm=[0, 3, 1, 2])

 # Upsample by inserting zeros.
 x = tf.reshape(x, [num_channels, batch_size, in_height, 1, in_width, 1])
 x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, upx - 1], [0, 0], [0, upy - 1]])
 x = tf.reshape(x, [batch_size, num_channels, in_height * upy, in_width * upx])

 # Pad or crop.
 x = tf.pad(x, [[0, 0], [0, 0],
 [tf.math.maximum(padx0, 0), tf.math.maximum(padx1, 0)],
 [tf.math.maximum(pady0, 0), tf.math.maximum(pady1, 0)]])
 x = x[:, :,
 tf.math.maximum(-pady0, 0) : x.shape[2] - tf.math.maximum(-pady1, 0),
 tf.math.maximum(-padx0, 0) : x.shape[3] - tf.math.maximum(-padx1, 0)]

 # Setup filter.
 f = f * (gain ** (f.ndim / 2))
 f = tf.cast(f, dtype=x.dtype)
 if not flip_filter:
 f = tf.reverse(f, axis=[-1])
 f = tf.reshape(f, shape=(1, 1, f.shape[-1]))
 f = tf.repeat(f, repeats=num_channels, axis=0)

 if tf.rank(f) == 4:
 f_0 = tf.transpose(f, perm=[2, 3, 1, 0])
 x = tf.nn.conv2d(x, f_0, 1, 'VALID')
 else:
 f_0 = tf.expand_dims(f, axis=2)
 f_0 = tf.transpose(f_0, perm=[2, 3, 1, 0])

 f_1 = tf.expand_dims(f, axis=3)
 f_1 = tf.transpose(f_1, perm=[2, 3, 1, 0])

 x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 x = tf.nn.conv2d(x, f_1, 1, 'VALID', data_format='NCHW')

 x = x[:, :, ::downy, ::downx]

 # Back to channel last.
 x = tf.transpose(x, perm=[0, 2, 3, 1])
 return x


 def filtered_lrelu(self,
 x, fu=None, fd=None, b=None,
 up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):
 #fu_w, fu_h = self.get_filter_size(fu)
 #fd_w, fd_h = self.get_filter_size(fd)

 px0, px1, py0, py1 = self.parse_padding(padding)

 #batch_size, in_h, in_w, channels = x.shape
 #in_dtype = x.dtype
 #out_w = (in_w * up + (px0 + px1) - (fu_w - 1) - (fd_w - 1) + (down - 1)) // down
 #out_h = (in_h * up + (py0 + py1) - (fu_h - 1) - (fd_h - 1) + (down - 1)) // down

 x = self.bias_act(x=x, b=b)
 x = self.upfirdn2d(x, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter)
 x = self.bias_act(x=x, act='lrelu', alpha=slope, gain=gain, clamp=clamp)
 x = self.upfirdn2d(x=x, f=fd, down=down, flip_filter=flip_filter)

 return x


 def call(self, inputs):
 return self.filtered_lrelu(inputs,
 fu=self.u_filter,
 fd=self.d_filter,
 b=self.bias,
 up=self.u_factor,
 down=self.d_factor,
 padding=self.padding,
 gain=self.gain,
 slope=self.slope,
 clamp=self.conv_clamp)

 def get_config(self):
 base_config = super(FilteredReLU, self).get_config()
 return base_config


class SynthesisLayer(Layer):

 def __init__(self,
 critically_sampled,

 in_channels,
 out_channels,
 in_size,
 out_size,
 in_sampling_rate,
 out_sampling_rate,
 in_cutoff,
 out_cutoff,
 in_half_width,
 out_half_width,

 conv_kernel = 3,
 lrelu_upsampling = 2,
 filter_size = 6,
 conv_clamp = 256,
 use_radial_filters = False,
 is_torgb = False,
 **kwargs):
 super(SynthesisLayer, self).__init__(**kwargs)
 self.critically_sampled = critically_sampled

 self.in_channels = in_channels
 self.out_channels = out_channels
 self.in_size = np.broadcast_to(np.asarray(in_size), [2])
 self.out_size = np.broadcast_to(np.asarray(out_size), [2])
 self.in_sampling_rate = in_sampling_rate
 self.out_sampling_rate = out_sampling_rate
 self.in_cutoff = in_cutoff
 self.out_cutoff = out_cutoff
 self.in_half_width = in_half_width
 self.out_half_width = out_half_width

 self.is_torgb = is_torgb

 self.conv_kernel = 1 if is_torgb else conv_kernel
 self.lrelu_upsampling = lrelu_upsampling
 self.conv_clamp = conv_clamp

 self.tmp_sampling_rate = max(in_sampling_rate, out_sampling_rate) * (1 if is_torgb else lrelu_upsampling)

 # Up sampling filter
 self.u_factor = int(np.rint(self.tmp_sampling_rate / self.in_sampling_rate))
 assert self.in_sampling_rate * self.u_factor == self.tmp_sampling_rate
 self.u_taps = filter_size * self.u_factor if self.u_factor > 1 and not self.is_torgb else 1
 self.u_filter = self.design_lowpass_filter(numtaps=self.u_taps,
 cutoff=self.in_cutoff,
 width=self.in_half_width*2,
 fs=self.tmp_sampling_rate)

 # Down sampling filter
 self.d_factor = int(np.rint(self.tmp_sampling_rate / self.out_sampling_rate))
 assert self.out_sampling_rate * self.d_factor == self.tmp_sampling_rate
 self.d_taps = filter_size * self.d_factor if self.d_factor > 1 and not self.is_torgb else 1
 self.d_radial = use_radial_filters and not self.critically_sampled
 self.d_filter = self.design_lowpass_filter(numtaps=self.d_taps,
 cutoff=self.out_cutoff,
 width=self.out_half_width*2,
 fs=self.tmp_sampling_rate,
 radial=self.d_radial)
 # Compute padding
 pad_total = (self.out_size - 1) * self.d_factor + 1
 pad_total -= (self.in_size + self.conv_kernel - 1) * self.u_factor
 pad_total += self.u_taps + self.d_taps - 2
 pad_lo = (pad_total + self.u_factor) // 2
 pad_hi = pad_total - pad_lo
 self.padding = [int(pad_lo[0]), int(pad_hi[0]), int(pad_lo[1]), int(pad_hi[1])]

 self.gain = 1 if self.is_torgb else np.sqrt(2)
 self.slope = 1 if self.is_torgb else 0.2

 self.act_funcs = {'linear':
 {'func': lambda x, **_: x,
 'def_alpha': 0,
 'def_gain': 1},
 'lrelu':
 {'func': lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha),
 'def_alpha': 0.2,
 'def_gain': np.sqrt(2)},
 }

 b_init = tf.zeros_initializer()
 self.bias = tf.Variable(initial_value=b_init(shape=(out_channels,),
 dtype="float32"),
 trainable=True)
 self.affine = Dense(self.in_channels)
 self.conv = Conv2DMod(self.out_channels, kernel_size=self.conv_kernel, padding='same')

 def design_lowpass_filter(self, numtaps, cutoff, width, fs, radial=False):
 if numtaps == 1:
 return None

 if not radial:
 f = sps.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs)
 return f

 x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
 r = np.hypot(*np.meshgrid(x, x))
 f = spspec.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
 beta = sps.kaiser_beta(sps.kaiser_atten(numtaps, width / (fs / 2)))
 w = np.kaiser(numtaps, beta)
 f *= np.outer(w, w)
 f /= np.sum(f)
 return f

 def get_filter_size(self, f):
 if f is None:
 return 1, 1
 assert 1 <= f.ndim <= 2
 return f.shape[-1], f.shape[0] # width, height

 def parse_padding(self, padding):
 if isinstance(padding, int):
 padding = [padding, padding]
 assert isinstance(padding, (list, tuple))
 assert all(isinstance(x, (int, np.integer)) for x in padding)
 padding = [int(x) for x in padding]
 if len(padding) == 2:
 px, py = padding
 padding = [px, px, py, py]
 px0, px1, py0, py1 = padding
 return px0, px1, py0, py1

 def bias_act(self, x, b=None, dim=3, act='linear', alpha=None, gain=None, clamp=None):
 spec = self.act_funcs[act]
 alpha = float(alpha if alpha is not None else spec['def_alpha'])
 gain = float(gain if gain is not None else spec['def_gain'])
 clamp = float(clamp if clamp is not None else -1)

 if b is not None:
 x = x + tf.reshape(b, shape=[-1 if i == dim else 1 for i in range(len(x.shape))])
 x = spec['func'](x, alpha=alpha)

 if gain != 1:
 x = x * gain

 if clamp >= 0:
 x = tf.clip_by_value(x, -clamp, clamp)
 return x

 def parse_scaling(self, scaling):
 if isinstance(scaling, int):
 scaling = [scaling, scaling]
 sx, sy = scaling
 assert sx >= 1 and sy >= 1
 return sx, sy

 def upfirdn2d(self, x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
 if f is None:
 f = tf.ones([1, 1], dtype=tf.float32)

 batch_size, in_height, in_width, num_channels = x.shape

 upx, upy = self.parse_scaling(up)
 downx, downy = self.parse_scaling(down)
 padx0, padx1, pady0, pady1 = self.parse_padding(padding)

 upW = in_width * upx + padx0 + padx1
 upH = in_height * upy + pady0 + pady1
 assert upW >= f.shape[-1] and upH >= f.shape[0]

 # Channel first format.
 x = tf.transpose(x, perm=[0, 3, 1, 2])

 # Upsample by inserting zeros.
 x = tf.reshape(x, [num_channels, batch_size, in_height, 1, in_width, 1])
 x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, upx - 1], [0, 0], [0, upy - 1]])
 x = tf.reshape(x, [batch_size, num_channels, in_height * upy, in_width * upx])

 # Pad or crop.
 x = tf.pad(x, [[0, 0], [0, 0],
 [tf.math.maximum(padx0, 0), tf.math.maximum(padx1, 0)],
 [tf.math.maximum(pady0, 0), tf.math.maximum(pady1, 0)]])
 x = x[:, :, tf.math.maximum(-pady0, 0) : x.shape[2] - tf.math.maximum(-pady1, 0), tf.math.maximum(-padx0, 0) : x.shape[3] - tf.math.maximum(-padx1, 0)]

 # Setup filter.
 f = f * (gain ** (f.ndim / 2))
 f = tf.cast(f, dtype=x.dtype)
 if not flip_filter:
 f = tf.reverse(f, axis=[-1])
 f = tf.reshape(f, shape=(1, 1, f.shape[-1]))
 f = tf.repeat(f, repeats=num_channels, axis=0)

 if tf.rank(f) == 4:
 f_0 = tf.transpose(f, perm=[2, 3, 1, 0])
 x = tf.nn.conv2d(x, f_0, 1, 'VALID')
 else:
 f_0 = tf.expand_dims(f, axis=2)
 f_0 = tf.transpose(f_0, perm=[2, 3, 1, 0])

 f_1 = tf.expand_dims(f, axis=3)
 f_1 = tf.transpose(f_1, perm=[2, 3, 1, 0])

 x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 x = tf.nn.conv2d(x, f_1, 1, 'VALID', data_format='NCHW')

 x = x[:, :, ::downy, ::downx]

 # Back to channel last.
 x = tf.transpose(x, perm=[0, 2, 3, 1])
 return x


 def filtered_lrelu(self,
 x, fu=None, fd=None, b=None,
 up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):
 #fu_w, fu_h = self.get_filter_size(fu)
 #fd_w, fd_h = self.get_filter_size(fd)

 px0, px1, py0, py1 = self.parse_padding(padding)

 #batch_size, in_h, in_w, channels = x.shape
 #in_dtype = x.dtype
 #out_w = (in_w * up + (px0 + px1) - (fu_w - 1) - (fd_w - 1) + (down - 1)) // down
 #out_h = (in_h * up + (py0 + py1) - (fu_h - 1) - (fd_h - 1) + (down - 1)) // down

 x = self.bias_act(x=x, b=b)
 x = self.upfirdn2d(x, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter)
 x = self.bias_act(x=x, act='lrelu', alpha=slope, gain=gain, clamp=clamp)
 x = self.upfirdn2d(x=x, f=fd, down=down, flip_filter=flip_filter)

 return x


 def call(self, inputs):
 x, w = inputs
 styles = self.affine(w)
 x = self.conv([x, styles])
 x = self.filtered_lrelu(x,
 fu=self.u_filter,
 fd=self.d_filter,
 b=self.bias,
 up=self.u_factor,
 down=self.d_factor,
 padding=self.padding,
 gain=self.gain,
 slope=self.slope,
 clamp=self.conv_clamp)
 return x

 def get_config(self):
 base_config = super(SynthesisLayer, self).get_config()
 return base_config


class SynthesisLayerNoMod(Layer):

 def __init__(self,
 critically_sampled,

 in_channels,
 out_channels,
 in_size,
 out_size,
 in_sampling_rate,
 out_sampling_rate,
 in_cutoff,
 out_cutoff,
 in_half_width,
 out_half_width,

 conv_kernel = 3,
 lrelu_upsampling = 2,
 filter_size = 6,
 conv_clamp = 256,
 use_radial_filters = False,
 is_torgb = False,
 batch_size = 10,
 **kwargs):
 super(SynthesisLayerNoMod, self).__init__(**kwargs)
 self.critically_sampled = critically_sampled
 self.bs = batch_size

 self.in_channels = in_channels
 self.out_channels = out_channels
 self.in_size = np.broadcast_to(np.asarray(in_size), [2])
 self.out_size = np.broadcast_to(np.asarray(out_size), [2])
 self.in_sampling_rate = in_sampling_rate
 self.out_sampling_rate = out_sampling_rate
 self.in_cutoff = in_cutoff
 self.out_cutoff = out_cutoff
 self.in_half_width = in_half_width
 self.out_half_width = out_half_width

 self.is_torgb = is_torgb

 self.conv_kernel = 1 if is_torgb else conv_kernel
 self.lrelu_upsampling = lrelu_upsampling
 self.conv_clamp = conv_clamp

 self.tmp_sampling_rate = max(in_sampling_rate, out_sampling_rate) * (1 if is_torgb else lrelu_upsampling)

 # Up sampling filter
 self.u_factor = int(np.rint(self.tmp_sampling_rate / self.in_sampling_rate))
 assert self.in_sampling_rate * self.u_factor == self.tmp_sampling_rate
 self.u_taps = filter_size * self.u_factor if self.u_factor > 1 and not self.is_torgb else 1
 self.u_filter = self.design_lowpass_filter(numtaps=self.u_taps,
 cutoff=self.in_cutoff,
 width=self.in_half_width*2,
 fs=self.tmp_sampling_rate)

 # Down sampling filter
 self.d_factor = int(np.rint(self.tmp_sampling_rate / self.out_sampling_rate))
 assert self.out_sampling_rate * self.d_factor == self.tmp_sampling_rate
 self.d_taps = filter_size * self.d_factor if self.d_factor > 1 and not self.is_torgb else 1
 self.d_radial = use_radial_filters and not self.critically_sampled
 self.d_filter = self.design_lowpass_filter(numtaps=self.d_taps,
 cutoff=self.out_cutoff,
 width=self.out_half_width*2,
 fs=self.tmp_sampling_rate,
 radial=self.d_radial)
 # Compute padding
 pad_total = (self.out_size - 1) * self.d_factor + 1
 pad_total -= (self.in_size + self.conv_kernel - 1) * self.u_factor
 pad_total += self.u_taps + self.d_taps - 2
 pad_lo = (pad_total + self.u_factor) // 2
 pad_hi = pad_total - pad_lo
 self.padding = [int(pad_lo[0]), int(pad_hi[0]), int(pad_lo[1]), int(pad_hi[1])]

 self.gain = 1 if self.is_torgb else np.sqrt(2)
 self.slope = 1 if self.is_torgb else 0.2

 self.act_funcs = {'linear':
 {'func': lambda x, **_: x,
 'def_alpha': 0,
 'def_gain': 1},
 'lrelu':
 {'func': lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha),
 'def_alpha': 0.2,
 'def_gain': np.sqrt(2)},
 }

 b_init = tf.zeros_initializer()
 self.bias = tf.Variable(initial_value=b_init(shape=(out_channels,),
 dtype="float32"),
 trainable=True)
 self.conv = Conv2D(self.out_channels, kernel_size=self.conv_kernel, padding='same')

 def design_lowpass_filter(self, numtaps, cutoff, width, fs, radial=False):
 if numtaps == 1:
 return None

 if not radial:
 f = sps.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs)
 return f

 x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
 r = np.hypot(*np.meshgrid(x, x))
 f = spspec.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
 beta = sps.kaiser_beta(sps.kaiser_atten(numtaps, width / (fs / 2)))
 w = np.kaiser(numtaps, beta)
 f *= np.outer(w, w)
 f /= np.sum(f)
 return f

 def get_filter_size(self, f):
 if f is None:
 return 1, 1
 assert 1 <= f.ndim <= 2
 return f.shape[-1], f.shape[0] # width, height

 def parse_padding(self, padding):
 if isinstance(padding, int):
 padding = [padding, padding]
 assert isinstance(padding, (list, tuple))
 assert all(isinstance(x, (int, np.integer)) for x in padding)
 padding = [int(x) for x in padding]
 if len(padding) == 2:
 px, py = padding
 padding = [px, px, py, py]
 px0, px1, py0, py1 = padding
 return px0, px1, py0, py1

 @tf.function
 def bias_act(self, x, b=None, dim=3, act='linear', alpha=None, gain=None, clamp=None):
 spec = self.act_funcs[act]
 alpha = float(alpha if alpha is not None else spec['def_alpha'])
 gain = float(gain if gain is not None else spec['def_gain'])
 clamp = float(clamp if clamp is not None else -1)

 if b is not None:
 x = x + tf.reshape(b, shape=[-1 if i == dim else 1 for i in range(len(x.shape))])
 x = spec['func'](x, alpha=alpha)

 if gain != 1:
 x = x * gain

 if clamp >= 0:
 x = tf.clip_by_value(x, -clamp, clamp)
 return x

 def parse_scaling(self, scaling):
 if isinstance(scaling, int):
 scaling = [scaling, scaling]
 sx, sy = scaling
 assert sx >= 1 and sy >= 1
 return sx, sy

 @tf.function
 def upfirdn2d(self, x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
 if f is None:
 f = tf.ones([1, 1], dtype=tf.float32)

 batch_size, in_height, in_width, num_channels = x.shape
 batch_size = tf.shape(x)[0]

 upx, upy = self.parse_scaling(up)
 downx, downy = self.parse_scaling(down)
 padx0, padx1, pady0, pady1 = self.parse_padding(padding)

 upW = in_width * upx + padx0 + padx1
 upH = in_height * upy + pady0 + pady1
 assert upW >= f.shape[-1] and upH >= f.shape[0]

 # Channel first format.
 x = tf.transpose(x, perm=[0, 3, 1, 2])
 # Upsample by inserting zeros.
 x = tf.reshape(x, [batch_size, num_channels, in_height, 1, in_width, 1])
 x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, upx - 1], [0, 0], [0, upy - 1]])
 x = tf.reshape(x, [batch_size, num_channels, in_height * upy, in_width * upx])

 # Pad or crop.
 x = tf.pad(x, [[0, 0], [0, 0],
 [tf.math.maximum(padx0, 0), tf.math.maximum(padx1, 0)],
 [tf.math.maximum(pady0, 0), tf.math.maximum(pady1, 0)]])
 x = x[:, :,
 tf.math.maximum(-pady0, 0) : x.shape[2] - tf.math.maximum(-pady1, 0),
 tf.math.maximum(-padx0, 0) : x.shape[3] - tf.math.maximum(-padx1, 0)]

 # Setup filter.
 f = f * (gain ** (tf.rank(f) / 2))
 f = tf.cast(f, dtype=x.dtype)
 if not flip_filter:
 f = tf.reverse(f, axis=[-1])
 f = tf.reshape(f, shape=(1, 1, f.shape[-1]))
 f = tf.repeat(f, repeats=num_channels, axis=0)

 #if tf.rank(f) == 500:
 # f_0 = tf.transpose(f, perm=[2, 3, 1, 0])
 # x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 #else:
 f_0 = tf.expand_dims(f, axis=2)
 f_0 = tf.transpose(f_0, perm=[2, 3, 1, 0])

 f_1 = tf.expand_dims(f, axis=3)
 f_1 = tf.transpose(f_1, perm=[2, 3, 1, 0])

 x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 x = tf.nn.conv2d(x, f_1, 1, 'VALID', data_format='NCHW')

 x = x[:, :, ::downy, ::downx]

 # Back to channel last.
 x = tf.transpose(x, perm=[0, 2, 3, 1])
 return x

 @tf.function
 def filtered_lrelu(self,
 x, fu=None, fd=None, b=None,
 up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):
 #fu_w, fu_h = self.get_filter_size(fu)
 #fd_w, fd_h = self.get_filter_size(fd)

 px0, px1, py0, py1 = self.parse_padding(padding)

 #batch_size, in_h, in_w, channels = x.shape
 #in_dtype = x.dtype
 #out_w = (in_w * up + (px0 + px1) - (fu_w - 1) - (fd_w - 1) + (down - 1)) // down
 #out_h = (in_h * up + (py0 + py1) - (fu_h - 1) - (fd_h - 1) + (down - 1)) // down

 x = self.bias_act(x=x, b=b)
 x = self.upfirdn2d(x, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter)
 x = self.bias_act(x=x, act='lrelu', alpha=slope, gain=gain, clamp=clamp)
 x = self.upfirdn2d(x=x, f=fd, down=down, flip_filter=flip_filter)

 return x


 def call(self, inputs):
 x = inputs
 x = self.conv(x)
 x = self.filtered_lrelu(x,
 fu=self.u_filter,
 fd=self.d_filter,
 b=self.bias,
 up=self.u_factor,
 down=self.d_factor,
 padding=self.padding,
 gain=self.gain,
 slope=self.slope,
 clamp=self.conv_clamp)
 return x

 def get_config(self):
 base_config = super(SynthesisLayer, self).get_config()
 return base_config


class SynthesisLayerNoModBN(Layer):

 def __init__(self,
 critically_sampled,

 in_channels,
 out_channels,
 in_size,
 out_size,
 in_sampling_rate,
 out_sampling_rate,
 in_cutoff,
 out_cutoff,
 in_half_width,
 out_half_width,

 conv_kernel = 3,
 lrelu_upsampling = 2,
 filter_size = 6,
 conv_clamp = 256,
 use_radial_filters = False,
 is_torgb = False,
 batch_size = 10,
 **kwargs):
 super(SynthesisLayerNoModBN, self).__init__(**kwargs)
 self.critically_sampled = critically_sampled
 self.bs = batch_size

 self.in_channels = in_channels
 self.out_channels = out_channels
 self.in_size = np.broadcast_to(np.asarray(in_size), [2])
 self.out_size = np.broadcast_to(np.asarray(out_size), [2])
 self.in_sampling_rate = in_sampling_rate
 self.out_sampling_rate = out_sampling_rate
 self.in_cutoff = in_cutoff
 self.out_cutoff = out_cutoff
 self.in_half_width = in_half_width
 self.out_half_width = out_half_width

 self.is_torgb = is_torgb

 self.conv_kernel = 1 if is_torgb else conv_kernel
 self.lrelu_upsampling = lrelu_upsampling
 self.conv_clamp = conv_clamp

 self.tmp_sampling_rate = max(in_sampling_rate, out_sampling_rate) * (1.0 if is_torgb else lrelu_upsampling)

 # Up sampling filter
 self.u_factor = int(np.rint(self.tmp_sampling_rate / self.in_sampling_rate))
 assert self.in_sampling_rate * self.u_factor == self.tmp_sampling_rate
 self.u_taps = filter_size * self.u_factor if self.u_factor > 1 and not self.is_torgb else 1
 self.u_filter = self.design_lowpass_filter(numtaps=self.u_taps,
 cutoff=self.in_cutoff,
 width=self.in_half_width*2,
 fs=self.tmp_sampling_rate)

 # Down sampling filter
 self.d_factor = int(np.rint(self.tmp_sampling_rate / self.out_sampling_rate))
 assert self.out_sampling_rate * self.d_factor == self.tmp_sampling_rate
 self.d_taps = filter_size * self.d_factor if self.d_factor > 1 and not self.is_torgb else 1
 self.d_radial = use_radial_filters and not self.critically_sampled
 self.d_filter = self.design_lowpass_filter(numtaps=self.d_taps,
 cutoff=self.out_cutoff,
 width=self.out_half_width*2,
 fs=self.tmp_sampling_rate,
 radial=self.d_radial)
 # Compute padding
 pad_total = (self.out_size - 1) * self.d_factor + 1
 pad_total -= (self.in_size) * self.u_factor
 pad_total += self.u_taps + self.d_taps - 2
 pad_lo = (pad_total + self.u_factor) // 2
 pad_hi = pad_total - pad_lo
 self.padding = [int(pad_lo[0]), int(pad_hi[0]), int(pad_lo[1]), int(pad_hi[1])]

 self.gain = 1 if self.is_torgb else np.sqrt(2)
 self.slope = 1 if self.is_torgb else 0.2

 self.act_funcs = {'linear':
 {'func': lambda x, **_: x,
 'def_alpha': 0,
 'def_gain': 1},
 'lrelu':
 {'func': lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha),
 'def_alpha': 0.2,
 'def_gain': np.sqrt(2)},
 }

 b_init = tf.zeros_initializer()
 self.bias = tf.Variable(initial_value=b_init(shape=(out_channels,),
 dtype="float32"),
 trainable=True)
 self.conv = Conv2D(self.out_channels, kernel_size=self.conv_kernel, padding='same')
 self.bn = BatchNormalization()

 def design_lowpass_filter(self, numtaps, cutoff, width, fs, radial=False):
 if numtaps == 1:
 return None

 if not radial:
 f = sps.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs)
 return f

 x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
 r = np.hypot(*np.meshgrid(x, x))
 f = spspec.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
 beta = sps.kaiser_beta(sps.kaiser_atten(numtaps, width / (fs / 2)))
 w = np.kaiser(numtaps, beta)
 f *= np.outer(w, w)
 f /= np.sum(f)
 return f

 def get_filter_size(self, f):
 if f is None:
 return 1, 1
 assert 1 <= f.ndim <= 2
 return f.shape[-1], f.shape[0] # width, height

 def parse_padding(self, padding):
 if isinstance(padding, int):
 padding = [padding, padding]
 assert isinstance(padding, (list, tuple))
 assert all(isinstance(x, (int, np.integer)) for x in padding)
 padding = [int(x) for x in padding]
 if len(padding) == 2:
 px, py = padding
 padding = [px, px, py, py]
 px0, px1, py0, py1 = padding
 return px0, px1, py0, py1

 @tf.function
 def bias_act(self, x, b=None, dim=3, act='linear', alpha=None, gain=None, clamp=None):
 spec = self.act_funcs[act]
 alpha = float(alpha if alpha is not None else spec['def_alpha'])
 gain = tf.cast(gain if gain is not None else spec['def_gain'], tf.float32)
 clamp = float(clamp if clamp is not None else -1)

 if b is not None:
 x = x + tf.reshape(b, shape=[-1 if i == dim else 1 for i in range(len(x.shape))])
 x = spec['func'](x, alpha=alpha)

 x = x * gain

 if clamp >= 0:
 x = tf.clip_by_value(x, -clamp, clamp)
 return x

 def parse_scaling(self, scaling):
 if isinstance(scaling, int):
 scaling = [scaling, scaling]
 sx, sy = scaling
 assert sx >= 1 and sy >= 1
 return sx, sy

 @tf.function
 def upfirdn2d(self, x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
 if f is None:
 f = tf.ones([1, 1], dtype=tf.float32)

 batch_size, in_height, in_width, num_channels = x.shape
 batch_size = tf.shape(x)[0]

 upx, upy = self.parse_scaling(up)
 downx, downy = self.parse_scaling(down)
 padx0, padx1, pady0, pady1 = self.parse_padding(padding)

 upW = in_width * upx + padx0 + padx1
 upH = in_height * upy + pady0 + pady1
 assert upW >= f.shape[-1] and upH >= f.shape[0]

 # Channel first format.
 x = tf.transpose(x, perm=[0, 3, 1, 2])
 # Upsample by inserting zeros.
 x = tf.reshape(x, [batch_size, num_channels, in_height, 1, in_width, 1])
 x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, upx - 1], [0, 0], [0, upy - 1]])
 x = tf.reshape(x, [batch_size, num_channels, in_height * upy, in_width * upx])

 # Pad or crop.
 x = tf.pad(x, [[0, 0], [0, 0],
 [tf.math.maximum(padx0, 0), tf.math.maximum(padx1, 0)],
 [tf.math.maximum(pady0, 0), tf.math.maximum(pady1, 0)]])
 x = x[:, :,
 tf.math.maximum(-pady0, 0) : x.shape[2] - tf.math.maximum(-pady1, 0),
 tf.math.maximum(-padx0, 0) : x.shape[3] - tf.math.maximum(-padx1, 0)]

 # Setup filter.
 f = f * (gain ** (tf.rank(f) / 2))
 f = tf.cast(f, dtype=x.dtype)
 if not flip_filter:
 f = tf.reverse(f, axis=[-1])
 f = tf.reshape(f, shape=(1, 1, f.shape[-1]))
 f = tf.repeat(f, repeats=num_channels, axis=0)

 #if tf.rank(f) == 500:
 # f_0 = tf.transpose(f, perm=[2, 3, 1, 0])
 # x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 #else:
 f_0 = tf.expand_dims(f, axis=2)
 f_0 = tf.transpose(f_0, perm=[2, 3, 1, 0])

 f_1 = tf.expand_dims(f, axis=3)
 f_1 = tf.transpose(f_1, perm=[2, 3, 1, 0])

 x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 x = tf.nn.conv2d(x, f_1, 1, 'VALID', data_format='NCHW')

 x = x[:, :, ::downy, ::downx]

 # Back to channel last.
 x = tf.transpose(x, perm=[0, 2, 3, 1])
 return x

 @tf.function
 def filtered_lrelu(self,
 x, fu=None, fd=None, b=None,
 up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):
 #fu_w, fu_h = self.get_filter_size(fu)
 #fd_w, fd_h = self.get_filter_size(fd)

 px0, px1, py0, py1 = self.parse_padding(padding)

 #batch_size, in_h, in_w, channels = x.shape
 #in_dtype = x.dtype
 #out_w = (in_w * up + (px0 + px1) - (fu_w - 1) - (fd_w - 1) + (down - 1)) // down
 #out_h = (in_h * up + (py0 + py1) - (fu_h - 1) - (fd_h - 1) + (down - 1)) // down

 x = self.bias_act(x=x, b=b)
 x = self.upfirdn2d(x, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter)
 x = self.bias_act(x=x, act='lrelu', alpha=slope, gain=gain, clamp=clamp)
 x = self.upfirdn2d(x=x, f=fd, down=down, flip_filter=flip_filter)

 return x


 def call(self, inputs):
 x = inputs
 x = self.conv(x)
 x = self.bn(x)
 x = self.filtered_lrelu(x,
 fu=self.u_filter,
 fd=self.d_filter,
 b=self.bias,
 up=self.u_factor,
 down=self.d_factor,
 padding=self.padding,
 gain=self.gain,
 slope=self.slope,
 clamp=self.conv_clamp)
 return x

 def get_config(self):
 base_config = super(SynthesisLayerNoModBN, self).get_config()
 return base_config


class SynthesisInput(Layer):
 def __init__(self,
 w_dim,
 channels,
 size,
 sampling_rate,
 bandwidth,
 **kwargs):
 super(SynthesisInput, self).__init__(**kwargs)
 self.w_dim = w_dim
 self.channels = channels
 self.size = np.broadcast_to(np.asarray(size), [2])
 self.sampling_rate = sampling_rate
 self.bandwidth = bandwidth

 # Draw random frequencies from uniform 2D disc.
 freqs = np.random.normal(size=(int(channels[0]), 2))
 radii = np.sqrt(np.sum(np.square(freqs), axis=1, keepdims=True))
 freqs /= radii * np.power(np.exp(np.square(radii)), 0.25)
 freqs *= bandwidth
 phases = np.random.uniform(size=[int(channels[0])]) - 0.5

 # Setup parameters and buffers.
 w_init = tf.random_normal_initializer()
 self.weight = tf.Variable(initial_value=w_init(shape=(self.channels, self.channels),
 dtype="float32"), rainable=True)
 self.affine = Dense(4, kernel_initializer=tf.zeros_initializer, bias_initializer=tf.zeros_initializer)
 self.transform = tf.eye(3, 3)
 self.freqs = tf.constant(freqs)
 self.phases = tf.constant(phases)

 def call(self, w):
 # Batch dimension
 transforms = tf.expand_dims(self.transform, axis=0)
 freqs = tf.expand_dims(self.freqs, axis=0)
 phases = tf.expand_dims(self.phases, axis=0)

 # Apply learned transformation.
 t = self.affine(w) # t = (r_c, r_s, t_x, t_y)
 t = t / tf.linalg.norm(t[:, :2], axis=1, keepdims=True)
 # Inverse rotation wrt. resulting image.
 m_r = tf.repeat(tf.expand_dims(tf.eye(3), axis=0), repeats=w.shape[0], axis=0)
 m_r[:, 0, 0] = t[:, 0] # r'_c
 m_r[:, 0, 1] = -t[:, 1] # r'_s
 m_r[:, 1, 0] = t[:, 1] # r'_s
 m_r[:, 1, 1] = t[:, 0] # r'_c
 # Inverse translation wrt. resulting image.
 m_t = tf.repeat(tf.expand_dims(tf.eye(3), axis=0), repeats=w.shape[0], axis=0)
 m_t[:, 0, 2] = -t[:, 2] # t'_x
 m_t[:, 1, 2] = -t[:, 3] # t'_y
 transforms = m_r @ m_t @ transforms

 # Transform frequencies.
 phases = phases + tf.expand_dims(freqs @ transforms[:, :2, 2:], axis=2)
 freqs = freqs @ transforms[:, :2, :2]

 # Dampen out-of-band frequencies that may occur due to the user-specified transform.
 amplitudes = tf.clip_by_value(1 - (tf.linalg.norm(freqs, axis=1, keepdims=True) - self.bandwidth) / (self.sampling_rate / 2 - self.bandwidth), 0, 1)



 def get_config(self):
 base_config = super(SynthesisInput, self).get_config()
 return base_config


class SynthesisLayerFS(Layer):

 def __init__(self,
 critically_sampled,

 in_channels,
 out_channels,
 in_size,
 out_size,
 in_sampling_rate,
 out_sampling_rate,
 in_cutoff,
 out_cutoff,
 in_half_width,
 out_half_width,

 conv_kernel = 3,
 lrelu_upsampling = 2,
 filter_size = 6,
 conv_clamp = 256,
 use_radial_filters = False,
 is_torgb = False,
 **kwargs):
 super(SynthesisLayerFS, self).__init__(**kwargs)
 self.critically_sampled = critically_sampled

 self.in_channels = in_channels
 self.out_channels = out_channels
 self.in_size = np.broadcast_to(np.asarray(in_size), [2])
 self.out_size = np.broadcast_to(np.asarray(out_size), [2])
 self.in_sampling_rate = in_sampling_rate
 self.out_sampling_rate = out_sampling_rate
 self.in_cutoff = in_cutoff
 self.out_cutoff = out_cutoff
 self.in_half_width = in_half_width
 self.out_half_width = out_half_width

 self.is_torgb = is_torgb

 self.conv_kernel = 1 if is_torgb else conv_kernel
 self.lrelu_upsampling = lrelu_upsampling
 self.conv_clamp = conv_clamp

 self.tmp_sampling_rate = max(in_sampling_rate, out_sampling_rate) * (1 if is_torgb else lrelu_upsampling)

 # Up sampling filter
 self.u_factor = int(np.rint(self.tmp_sampling_rate / self.in_sampling_rate))
 assert self.in_sampling_rate * self.u_factor == self.tmp_sampling_rate
 self.u_taps = filter_size * self.u_factor if self.u_factor > 1 and not self.is_torgb else 1
 self.u_filter = self.design_lowpass_filter(numtaps=self.u_taps,
 cutoff=self.in_cutoff,
 width=self.in_half_width*2,
 fs=self.tmp_sampling_rate)

 # Down sampling filter
 self.d_factor = int(np.rint(self.tmp_sampling_rate / self.out_sampling_rate))
 assert self.out_sampling_rate * self.d_factor == self.tmp_sampling_rate
 self.d_taps = filter_size * self.d_factor if self.d_factor > 1 and not self.is_torgb else 1
 self.d_radial = use_radial_filters and not self.critically_sampled
 self.d_filter = self.design_lowpass_filter(numtaps=self.d_taps,
 cutoff=self.out_cutoff,
 width=self.out_half_width*2,
 fs=self.tmp_sampling_rate,
 radial=self.d_radial)
 # Compute padding
 pad_total = (self.out_size - 1) * self.d_factor + 1
 pad_total -= (self.in_size + self.conv_kernel - 1) * self.u_factor
 pad_total += self.u_taps + self.d_taps - 2
 pad_lo = (pad_total + self.u_factor) // 2
 pad_hi = pad_total - pad_lo
 self.padding = [int(pad_lo[0]), int(pad_hi[0]), int(pad_lo[1]), int(pad_hi[1])]

 self.gain = 1 if self.is_torgb else np.sqrt(2)
 self.slope = 1 if self.is_torgb else 0.2

 self.act_funcs = {'linear':
 {'func': lambda x, **_: x,
 'def_alpha': 0,
 'def_gain': 1},
 'lrelu':
 {'func': lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha),
 'def_alpha': 0.2,
 'def_gain': np.sqrt(2)},
 }

 b_init = tf.zeros_initializer()
 self.bias = tf.Variable(initial_value=b_init(shape=(out_channels,),
 dtype="float32"),
 trainable=True)
 self.affine = Dense(self.out_channels)
 self.conv_mod = Conv2DMod(self.out_channels, kernel_size=self.conv_kernel, padding='same')
 self.bn = BatchNormalization()
 self.conv_gamma = Conv2D(self.out_channels, kernel_size=1)
 self.conv_beta = Conv2D(self.out_channels, kernel_size=1)
 self.conv_gate = Conv2D(self.out_channels, kernel_size=1)
 self.conv_final = Conv2D(self.out_channels, kernel_size=self.conv_kernel, padding='same')


 def design_lowpass_filter(self, numtaps, cutoff, width, fs, radial=False):
 if numtaps == 1:
 return None

 if not radial:
 f = sps.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs)
 return f

 x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
 r = np.hypot(*np.meshgrid(x, x))
 f = spspec.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
 beta = sps.kaiser_beta(sps.kaiser_atten(numtaps, width / (fs / 2)))
 w = np.kaiser(numtaps, beta)
 f *= np.outer(w, w)
 f /= np.sum(f)
 return f

 def get_filter_size(self, f):
 if f is None:
 return 1, 1
 assert 1 <= f.ndim <= 2
 return f.shape[-1], f.shape[0] # width, height

 def parse_padding(self, padding):
 if isinstance(padding, int):
 padding = [padding, padding]
 assert isinstance(padding, (list, tuple))
 assert all(isinstance(x, (int, np.integer)) for x in padding)
 padding = [int(x) for x in padding]
 if len(padding) == 2:
 px, py = padding
 padding = [px, px, py, py]
 px0, px1, py0, py1 = padding
 return px0, px1, py0, py1

 @tf.function
 def bias_act(self, x, b=None, dim=3, act='linear', alpha=None, gain=None, clamp=None):
 spec = self.act_funcs[act]
 alpha = float(alpha if alpha is not None else spec['def_alpha'])
 gain = tf.cast(gain if gain is not None else spec['def_gain'], tf.float32)
 clamp = float(clamp if clamp is not None else -1)

 if b is not None:
 x = x + tf.reshape(b, shape=[-1 if i == dim else 1 for i in range(len(x.shape))])
 x = spec['func'](x, alpha=alpha)

 x = x * gain

 if clamp >= 0:
 x = tf.clip_by_value(x, -clamp, clamp)
 return x

 def parse_scaling(self, scaling):
 if isinstance(scaling, int):
 scaling = [scaling, scaling]
 sx, sy = scaling
 assert sx >= 1 and sy >= 1
 return sx, sy

 @tf.function
 def upfirdn2d(self, x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
 if f is None:
 f = tf.ones([1, 1], dtype=tf.float32)

 batch_size, in_height, in_width, num_channels = x.shape
 batch_size = tf.shape(x)[0]

 upx, upy = self.parse_scaling(up)
 downx, downy = self.parse_scaling(down)
 padx0, padx1, pady0, pady1 = self.parse_padding(padding)

 upW = in_width * upx + padx0 + padx1
 upH = in_height * upy + pady0 + pady1
 assert upW >= f.shape[-1] and upH >= f.shape[0]

 # Channel first format.
 x = tf.transpose(x, perm=[0, 3, 1, 2])
 # Upsample by inserting zeros.
 x = tf.reshape(x, [batch_size, num_channels, in_height, 1, in_width, 1])
 x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, upx - 1], [0, 0], [0, upy - 1]])
 x = tf.reshape(x, [batch_size, num_channels, in_height * upy, in_width * upx])

 # Pad or crop.
 x = tf.pad(x, [[0, 0], [0, 0],
 [tf.math.maximum(padx0, 0), tf.math.maximum(padx1, 0)],
 [tf.math.maximum(pady0, 0), tf.math.maximum(pady1, 0)]])
 x = x[:, :,
 tf.math.maximum(-pady0, 0): x.shape[2] - tf.math.maximum(-pady1, 0),
 tf.math.maximum(-padx0, 0): x.shape[3] - tf.math.maximum(-padx1, 0)]

 # Setup filter.
 f = f * (gain ** (tf.rank(f) / 2))
 f = tf.cast(f, dtype=x.dtype)
 if not flip_filter:
 f = tf.reverse(f, axis=[-1])
 f = tf.reshape(f, shape=(1, 1, f.shape[-1]))
 f = tf.repeat(f, repeats=num_channels, axis=0)

 # if tf.rank(f) == 500:
 # f_0 = tf.transpose(f, perm=[2, 3, 1, 0])
 # x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 # else:
 f_0 = tf.expand_dims(f, axis=2)
 f_0 = tf.transpose(f_0, perm=[2, 3, 1, 0])

 f_1 = tf.expand_dims(f, axis=3)
 f_1 = tf.transpose(f_1, perm=[2, 3, 1, 0])

 x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 x = tf.nn.conv2d(x, f_1, 1, 'VALID', data_format='NCHW')

 x = x[:, :, ::downy, ::downx]

 # Back to channel last.
 x = tf.transpose(x, perm=[0, 2, 3, 1])
 return x

 @tf.function
 def filtered_lrelu(self,
 x, fu=None, fd=None, b=None,
 up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):
 #fu_w, fu_h = self.get_filter_size(fu)
 #fd_w, fd_h = self.get_filter_size(fd)

 px0, px1, py0, py1 = self.parse_padding(padding)

 #batch_size, in_h, in_w, channels = x.shape
 #in_dtype = x.dtype
 #out_w = (in_w * up + (px0 + px1) - (fu_w - 1) - (fd_w - 1) + (down - 1)) // down
 #out_h = (in_h * up + (py0 + py1) - (fu_h - 1) - (fd_h - 1) + (down - 1)) // down

 x = self.bias_act(x=x, b=b)
 x = self.upfirdn2d(x, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter)
 x = self.bias_act(x=x, act='lrelu', alpha=slope, gain=gain, clamp=clamp)
 x = self.upfirdn2d(x=x, f=fd, down=down, flip_filter=flip_filter)

 return x

 @tf.function
 def aadm(self, x, w, a):
 w_affine = self.affine(w)
 x_norm = self.bn(x)

 x_id = self.conv_mod([x_norm, w_affine])

 gate = self.conv_gate(x_norm)
 gate = tf.nn.sigmoid(gate)

 x_att_beta = self.conv_beta(a)
 x_att_gamma = self.conv_gamma(a)

 x_att = x_norm * x_att_beta + x_att_gamma

 h = x_id * gate + (1 - gate) * x_att

 return h


 def call(self, inputs):
 x, w, a = inputs
 x = self.conv_final(x)
 x = self.aadm(x, w, a)
 x = self.filtered_lrelu(x,
 fu=self.u_filter,
 fd=self.d_filter,
 b=self.bias,
 up=self.u_factor,
 down=self.d_factor,
 padding=self.padding,
 gain=self.gain,
 slope=self.slope,
 clamp=self.conv_clamp)
 return x

 def get_config(self):
 base_config = super(SynthesisLayerFS, self).get_config()
 return base_config


class SynthesisLayerUpDownOnly(Layer):

 def __init__(self,
 critically_sampled,

 in_channels,
 out_channels,
 in_size,
 out_size,
 in_sampling_rate,
 out_sampling_rate,
 in_cutoff,
 out_cutoff,
 in_half_width,
 out_half_width,

 conv_kernel = 3,
 lrelu_upsampling = 2,
 filter_size = 6,
 conv_clamp = 256,
 use_radial_filters = False,
 is_torgb = False,
 **kwargs):
 super(SynthesisLayerUpDownOnly, self).__init__(**kwargs)
 self.critically_sampled = critically_sampled

 self.in_channels = in_channels
 self.out_channels = out_channels
 self.in_size = np.broadcast_to(np.asarray(in_size), [2])
 self.out_size = np.broadcast_to(np.asarray(out_size), [2])
 self.in_sampling_rate = in_sampling_rate
 self.out_sampling_rate = out_sampling_rate
 self.in_cutoff = in_cutoff
 self.out_cutoff = out_cutoff
 self.in_half_width = in_half_width
 self.out_half_width = out_half_width

 self.is_torgb = is_torgb

 self.conv_kernel = 1 if is_torgb else conv_kernel
 self.lrelu_upsampling = lrelu_upsampling
 self.conv_clamp = conv_clamp

 self.tmp_sampling_rate = max(in_sampling_rate, out_sampling_rate) * (1 if is_torgb else lrelu_upsampling)

 # Up sampling filter
 self.u_factor = int(np.rint(self.tmp_sampling_rate / self.in_sampling_rate))
 assert self.in_sampling_rate * self.u_factor == self.tmp_sampling_rate
 self.u_taps = filter_size * self.u_factor if self.u_factor > 1 and not self.is_torgb else 1
 self.u_filter = self.design_lowpass_filter(numtaps=self.u_taps,
 cutoff=self.in_cutoff,
 width=self.in_half_width*2,
 fs=self.tmp_sampling_rate)

 # Down sampling filter
 self.d_factor = int(np.rint(self.tmp_sampling_rate / self.out_sampling_rate))
 assert self.out_sampling_rate * self.d_factor == self.tmp_sampling_rate
 self.d_taps = filter_size * self.d_factor if self.d_factor > 1 and not self.is_torgb else 1
 self.d_radial = use_radial_filters and not self.critically_sampled
 self.d_filter = self.design_lowpass_filter(numtaps=self.d_taps,
 cutoff=self.out_cutoff,
 width=self.out_half_width*2,
 fs=self.tmp_sampling_rate,
 radial=self.d_radial)
 # Compute padding
 pad_total = (self.out_size - 1) * self.d_factor + 1
 pad_total -= (self.in_size + self.conv_kernel - 1) * self.u_factor
 pad_total += self.u_taps + self.d_taps - 2
 pad_lo = (pad_total + self.u_factor) // 2
 pad_hi = pad_total - pad_lo
 self.padding = [int(pad_lo[0]), int(pad_hi[0]), int(pad_lo[1]), int(pad_hi[1])]

 self.gain = 1 if self.is_torgb else np.sqrt(2)
 self.slope = 1 if self.is_torgb else 0.2

 self.act_funcs = {'linear':
 {'func': lambda x, **_: x,
 'def_alpha': 0,
 'def_gain': 1},
 'lrelu':
 {'func': lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha),
 'def_alpha': 0.2,
 'def_gain': np.sqrt(2)},
 }

 def design_lowpass_filter(self, numtaps, cutoff, width, fs, radial=False):
 if numtaps == 1:
 return None

 if not radial:
 f = sps.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs)
 return f

 x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
 r = np.hypot(*np.meshgrid(x, x))
 f = spspec.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
 beta = sps.kaiser_beta(sps.kaiser_atten(numtaps, width / (fs / 2)))
 w = np.kaiser(numtaps, beta)
 f *= np.outer(w, w)
 f /= np.sum(f)
 return f

 def get_filter_size(self, f):
 if f is None:
 return 1, 1
 assert 1 <= f.ndim <= 2
 return f.shape[-1], f.shape[0] # width, height

 def parse_padding(self, padding):
 if isinstance(padding, int):
 padding = [padding, padding]
 assert isinstance(padding, (list, tuple))
 assert all(isinstance(x, (int, np.integer)) for x in padding)
 padding = [int(x) for x in padding]
 if len(padding) == 2:
 px, py = padding
 padding = [px, px, py, py]
 px0, px1, py0, py1 = padding
 return px0, px1, py0, py1

 def bias_act(self, x, b=None, dim=3, act='linear', alpha=None, gain=None, clamp=None):
 spec = self.act_funcs[act]
 alpha = float(alpha if alpha is not None else spec['def_alpha'])
 gain = float(gain if gain is not None else spec['def_gain'])
 clamp = float(clamp if clamp is not None else -1)

 if b is not None:
 x = x + tf.reshape(b, shape=[-1 if i == dim else 1 for i in range(len(x.shape))])
 x = spec['func'](x, alpha=alpha)

 if gain != 1:
 x = x * gain

 if clamp >= 0:
 x = tf.clip_by_value(x, -clamp, clamp)
 return x

 def parse_scaling(self, scaling):
 if isinstance(scaling, int):
 scaling = [scaling, scaling]
 sx, sy = scaling
 assert sx >= 1 and sy >= 1
 return sx, sy

 def upfirdn2d(self, x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
 if f is None:
 f = tf.ones([1, 1], dtype=tf.float32)

 batch_size, in_height, in_width, num_channels = x.shape

 upx, upy = self.parse_scaling(up)
 downx, downy = self.parse_scaling(down)
 padx0, padx1, pady0, pady1 = self.parse_padding(padding)

 upW = in_width * upx + padx0 + padx1
 upH = in_height * upy + pady0 + pady1
 assert upW >= f.shape[-1] and upH >= f.shape[0]

 # Channel first format.
 x = tf.transpose(x, perm=[0, 3, 1, 2])

 # Upsample by inserting zeros.
 x = tf.reshape(x, [num_channels, batch_size, in_height, 1, in_width, 1])
 x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, upx - 1], [0, 0], [0, upy - 1]])
 x = tf.reshape(x, [batch_size, num_channels, in_height * upy, in_width * upx])

 # Pad or crop.
 x = tf.pad(x, [[0, 0], [0, 0],
 [tf.math.maximum(padx0, 0), tf.math.maximum(padx1, 0)],
 [tf.math.maximum(pady0, 0), tf.math.maximum(pady1, 0)]])
 x = x[:, :, tf.math.maximum(-pady0, 0) : x.shape[2] - tf.math.maximum(-pady1, 0), tf.math.maximum(-padx0, 0) : x.shape[3] - tf.math.maximum(-padx1, 0)]

 # Setup filter.
 f = f * (gain ** (f.ndim / 2))
 f = tf.cast(f, dtype=x.dtype)
 if not flip_filter:
 f = tf.reverse(f, axis=[-1])
 f = tf.reshape(f, shape=(1, 1, f.shape[-1]))
 f = tf.repeat(f, repeats=num_channels, axis=0)

 if tf.rank(f) == 4:
 f_0 = tf.transpose(f, perm=[2, 3, 1, 0])
 x = tf.nn.conv2d(x, f_0, 1, 'VALID')
 else:
 f_0 = tf.expand_dims(f, axis=2)
 f_0 = tf.transpose(f_0, perm=[2, 3, 1, 0])

 f_1 = tf.expand_dims(f, axis=3)
 f_1 = tf.transpose(f_1, perm=[2, 3, 1, 0])

 x = tf.nn.conv2d(x, f_0, 1, 'VALID', data_format='NCHW')
 x = tf.nn.conv2d(x, f_1, 1, 'VALID', data_format='NCHW')

 x = x[:, :, ::downy, ::downx]

 # Back to channel last.
 x = tf.transpose(x, perm=[0, 2, 3, 1])
 return x


 def filtered_lrelu(self,
 x, fu=None, fd=None, b=None,
 up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):

 px0, px1, py0, py1 = self.parse_padding(padding)

 x = self.upfirdn2d(x, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter)
 x = self.bias_act(x=x, act='lrelu', alpha=slope, gain=gain, clamp=clamp)
 x = self.upfirdn2d(x=x, f=fd, down=down, flip_filter=flip_filter)

 return x


 def call(self, inputs):
 x = inputs
 x = self.filtered_lrelu(x,
 fu=self.u_filter,
 fd=self.d_filter,
 b=self.bias,
 up=self.u_factor,
 down=self.d_factor,
 padding=self.padding,
 gain=self.gain,
 slope=self.slope,
 clamp=self.conv_clamp)
 return x

 def get_config(self):
 base_config = super(SynthesisLayerUpDownOnly, self).get_config()
 return base_config


class Localization(Layer):
 def __init__(self):
 super(Localization, self).__init__()

 self.pool = MaxPooling2D()
 self.conv_0 = Conv2D(36, 5, activation='relu')
 self.conv_1 = Conv2D(36, 5, activation='relu')
 self.flatten = Flatten()
 self.fc_0 = Dense(36, activation='relu')
 self.fc_1 = Dense(6, bias_initializer=tf.keras.initializers.constant([1.0, 0.0, 0.0, 0.0, 1.0, 0.0]),
 kernel_initializer='zeros')
 self.reshape = Reshape((2, 3))

 def build(self, input_shape):
 print(input_shape)

 def compute_output_shape(self, input_shape):
 return [None, 6]

 def call(self, inputs):
 x = self.conv_0(inputs)
 x = self.pool(x)
 x = self.conv_1(x)
 x = self.pool(x)
 x = self.flatten(x)
 x = self.fc_0(x)
 theta = self.fc_1(x)
 theta = self.reshape(theta)

 return theta


class BilinearInterpolation(Layer):
 def __init__(self, height=36, width=36):
 super(BilinearInterpolation, self).__init__()
 self.height = height
 self.width = width

 def compute_output_shape(self, input_shape):
 return [None, self.height, self.width, 1]

 def get_config(self):
 config = {
 'height': self.height,
 'width': self.width
 }
 base_config = super(BilinearInterpolation, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))

 def advance_indexing(self, inputs, x, y):
 shape = tf.shape(inputs)
 batch_size = shape[0]

 batch_idx = tf.range(0, batch_size)
 batch_idx = tf.reshape(batch_idx, (batch_size, 1, 1))

 b = tf.tile(batch_idx, (1, self.height, self.width))
 indices = tf.stack([b, y, x], 3)

 return tf.gather_nd(inputs, indices)

 def grid_generator(self, batch):
 x = tf.linspace(-1, 1, self.width)
 y = tf.linspace(-1, 1, self.height)

 xx, yy = tf.meshgrid(x, y)
 xx = tf.reshape(xx, (-1,))
 yy = tf.reshape(yy, (-1,))

 homogenous_coordinates = tf.stack([xx, yy, tf.ones_like(xx)])
 homogenous_coordinates = tf.expand_dims(homogenous_coordinates, axis=0)
 homogenous_coordinates = tf.tile(homogenous_coordinates, [batch, 1, 1])
 homogenous_coordinates = tf.cast(homogenous_coordinates, dtype=tf.float32)
 return homogenous_coordinates

 def interpolate(self, images, homogenous_coordinates, theta):

 with tf.name_scope("Transformation"):
 transformed = tf.matmul(theta, homogenous_coordinates)
 transformed = tf.transpose(transformed, perm=[0, 2, 1])
 transformed = tf.reshape(transformed, [-1, self.height, self.width, 2])

 x_transformed = transformed[:, :, :, 0]
 y_transformed = transformed[:, :, :, 1]

 x = ((x_transformed + 1.) * tf.cast(self.width, dtype=tf.float32)) * 0.5
 y = ((y_transformed + 1.) * tf.cast(self.height, dtype=tf.float32)) * 0.5

 with tf.name_scope("VaribleCasting"):
 x0 = tf.cast(tf.math.floor(x), dtype=tf.int32)
 x1 = x0 + 1
 y0 = tf.cast(tf.math.floor(y), dtype=tf.int32)
 y1 = y0 + 1

 x0 = tf.clip_by_value(x0, 0, self.width-1)
 x1 = tf.clip_by_value(x1, 0, self.width - 1)
 y0 = tf.clip_by_value(y0, 0, self.height - 1)
 y1 = tf.clip_by_value(y1, 0, self.height - 1)
 x = tf.clip_by_value(x, 0, tf.cast(self.width, dtype=tf.float32) - 1.0)
 y = tf.clip_by_value(y, 0, tf.cast(self.height, dtype=tf.float32) - 1.0)

 with tf.name_scope("AdvancedIndexing"):
 i_a = self.advance_indexing(images, x0, y0)
 i_b = self.advance_indexing(images, x0, y1)
 i_c = self.advance_indexing(images, x1, y0)
 i_d = self.advance_indexing(images, x1, y1)

 with tf.name_scope("Interpolation"):
 x0 = tf.cast(x0, dtype=tf.float32)
 x1 = tf.cast(x1, dtype=tf.float32)
 y0 = tf.cast(y0, dtype=tf.float32)
 y1 = tf.cast(y1, dtype=tf.float32)

 w_a = (x1 - x) * (y1 - y)
 w_b = (x1 - x) * (y - y0)
 w_c = (x - x0) * (y1 - y)
 w_d = (x - x0) * (y - y0)

 w_a = tf.expand_dims(w_a, axis=3)
 w_b = tf.expand_dims(w_b, axis=3)
 w_c = tf.expand_dims(w_c, axis=3)
 w_d = tf.expand_dims(w_d, axis=3)

 return tf.math.add_n([w_a * i_a + w_b * i_b + w_c * i_c + w_d * i_d])

 def call(self, inputs):
 images, theta = inputs
 homogenous_coordinates = self.grid_generator(batch=tf.shape(images)[0])
 return self.interpolate(images, homogenous_coordinates, theta)


class ResBlockLR(Layer):
 def __init__(self, filters=16):
 super(ResBlockLR, self).__init__()
 self.filters = filters

 self.conv_0 = Conv2D(filters=filters,
 kernel_size=3,
 strides=1,
 padding='same')
 self.bn_0 = BatchNormalization()
 self.conv_1 = Conv2D(filters=filters,
 kernel_size=3,
 strides=1,
 padding='same')
 self.bn_1 = BatchNormalization()

 def get_config(self):
 config = {
 'filters': self.filters,
 }
 base_config = super(ResBlockLR, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))

 def call(self, inputs):
 x = self.conv_0(inputs)
 x = self.bn_0(x)
 x = tf.nn.leaky_relu(x, alpha=0.2)
 x = self.conv_1(x)
 x = self.bn_1(x)
 return x + inputs


class LearnedResize(Layer):
 def __init__(self, width, height, filters=16, in_channels=3, num_res_block=3, interpolation='bilinear'):
 super(LearnedResize, self).__init__()
 self.filters = filters
 self.num_res_block = num_res_block
 self.interpolation = interpolation
 self.in_channels = in_channels
 self.width = width
 self.height = height

 self.resize_layer = tf.keras.layers.experimental.preprocessing.Resizing(height,
 width,
 interpolation=interpolation)

 self.init_layers = tf.keras.models.Sequential([Conv2D(filters=filters,
 kernel_size=7,
 strides=1,
 padding='same'),
 LeakyReLU(0.2),
 Conv2D(filters=filters,
 kernel_size=1,
 strides=1,
 padding='same'),
 LeakyReLU(0.2),
 BatchNormalization()
 ])
 res_blocks = []
 for i in range(num_res_block):
 res_blocks.append(ResBlockLR(filters=filters))
 res_blocks.append(Conv2D(filters=filters,
 kernel_size=3,
 strides=1,
 padding='same',
 use_bias=False))
 res_blocks.append(BatchNormalization())
 self.res_block_pipe = tf.keras.models.Sequential(res_blocks)
 self.final_conv = Conv2D(filters=in_channels,
 kernel_size=3,
 strides=1,
 padding='same')


 def compute_output_shape(self, input_shape):
 return [None, self.target_size[0], self.target_size[1], input_shape[-1]]

 def get_config(self):
 config = {
 'filters': self.filters,
 'num_res_block': self.num_res_block,
 'interpolation': self.interpolation,
 'width': self.width,
 'height': self.height,
 }
 base_config = super(LearnedResize, self).get_config()
 return dict(list(base_config.items()) + list(config.items()))

 def call(self, inputs):
 x_l = self.init_layers(inputs)
 x_l_0 = self.resize_layer(x_l)
 x_l = self.res_block_pipe(x_l_0)
 x_l = x_l + x_l_0
 x_l = self.final_conv(x_l)

 x = self.resize_layer(inputs)

 return x + x_l