SAE: added new archi 'vg'

__dev/test.py

@@ -305,20 +305,98 @@ def get_transform_mat (image_landmarks, output_size, scale=1.0):
 #    alignments.append (dflimg.get_source_landmarks())
 import mathlib
 def main():
+
+    def f ( *args, asd=True, **kwargs ):
+        import code
+        code.interact(local=dict(globals(), **locals()))
+    
+    f( 1, asd=True, bg=0)
+    
     from nnlib import nnlib
-    exec( nnlib.import_all(), locals(), globals() )
+    exec( nnlib.import_all( device_config=nnlib.device.Config() ), locals(), globals() )
     PMLTile = nnlib.PMLTile
     PMLK = nnlib.PMLK
     
-    image = cv2.imread('D:\\DeepFaceLab\\test\\00000.png').astype(np.float32) / 255.0    
+    class DSSIMObjective:
-    image = cv2.resize ( image, (128,128) )
+        """Computes DSSIM index between img1 and img2.
+        This function is based on the standard SSIM implementation from:
+        Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004).
+        """
 
-    image = cv2.cvtColor (image, cv2.COLOR_BGR2GRAY)    
+        def __init__(self, k1=0.01, k2=0.03, max_value=1.0):
-    image = np.expand_dims (image, -1)
+            self.__name__ = 'DSSIMObjective'
+            self.k1 = k1
+            self.k2 = k2
+            self.max_value = max_value
+            self.c1 = (self.k1 * self.max_value) ** 2
+            self.c2 = (self.k2 * self.max_value) ** 2
+            self.dim_ordering = K.image_data_format()
+            self.backend = K.backend()
+
+        def __int_shape(self, x):
+            return K.int_shape(x) if self.backend == 'tensorflow' else K.shape(x)
+
+        def __call__(self, y_true, y_pred):
+            ch = K.shape(y_pred)[-1]
+
+            def softmax(x, axis=-1):
+                y = np.exp(x - np.max(x, axis, keepdims=True))
+                return y / np.sum(y, axis, keepdims=True)
+                    
+            def _fspecial_gauss(size, sigma):
+                #Function to mimic the 'fspecial' gaussian MATLAB function.
+                coords = np.arange(0, size, dtype=K.floatx())
+                coords -= (size - 1 ) / 2.0
+                g = coords**2
+                g *= ( -0.5 / (sigma**2) )
+                g = np.reshape (g, (1,-1)) + np.reshape(g, (-1,1) )
+                g = np.reshape (g, (1,-1))
+                g = softmax(g)
+                g = K.constant ( np.reshape (g, (size, size, 1, 1))  )
+                g = K.tile (g, (1,1,ch,1))
+                return g
+                      
+            kernel = _fspecial_gauss(11,1.5)
+
+            def reducer(x):
+                return K.depthwise_conv2d(x, kernel, strides=(1, 1), padding='valid')
+
+            c1 = (self.k1 * self.max_value) ** 2
+            c2 = (self.k2 * self.max_value) ** 2
+            
+            mean0 = reducer(y_true)
+            mean1 = reducer(y_pred)
+            num0 = mean0 * mean1 * 2.0
+            den0 = K.square(mean0) + K.square(mean1)
+            luminance = (num0 + c1) / (den0 + c1)
+            
+            num1 = reducer(y_true * y_pred) * 2.0
+            den1 = reducer(K.square(y_true) + K.square(y_pred))
+            c2 *= 1.0 #compensation factor
+            cs = (num1 - num0 + c2) / (den1 - den0 + c2)
+
+            ssim_val = K.mean(luminance * cs, axis=(-3, -2) )
+            return K.mean( (1.0 - ssim_val ) / 2.0 )
+
+    image = cv2.imread('D:\\DeepFaceLab\\test\\00000.png').astype(np.float32) / 255.0    
     image = np.expand_dims (image, 0)
     image_shape = image.shape
     
-    t = K.placeholder ( image_shape ) #K.constant ( np.ones ( (10,) ) )
+    image2 = cv2.imread('D:\\DeepFaceLab\\test\\00001.png').astype(np.float32) / 255.0    
+    image2 = np.expand_dims (image2, 0)
+    image2_shape = image2.shape
+    
+    #image = np.random.uniform ( size=(1,256,256,3) )
+    #image2 = np.random.uniform ( size=(1,256,256,3) )
+    
+    t1 = K.placeholder ( (None,) + image_shape[1:], name="t1" )
+    t2 = K.placeholder ( (None,) + image_shape[1:], name="t2" )
+    
+    l1_t = DSSIMObjective() (t1,t2 )
+    l1, = K.function([t1, t2],[l1_t]) ([image, image2])
+    
+    print (l1)
+    
     import code
     code.interact(local=dict(globals(), **locals()))
     
@@ -1279,4 +1357,34 @@ O[i0, i1, i2, i3: (1 + 1 - 1)/1, (64 + 1 - 1)/1, (64 + 2 - 1)/2, (1 + 1 - 1)/1]
 
 
 if __name__ == "__main__":
+    #os.environ["KERAS_BACKEND"] = "plaidml.keras.backend"
+    #os.environ["PLAIDML_DEVICE_IDS"] = "opencl_nvidia_geforce_gtx_1060_6gb.0"
+    #import keras
+    #K = keras.backend
+    #
+    #image = np.random.uniform ( size=(1,256,256,3) )
+    #image2 = np.random.uniform ( size=(1,256,256,3) )
+    #
+    #y_true = K.placeholder ( (None,) + image.shape[1:] )
+    #y_pred = K.placeholder ( (None,) + image2.shape[1:] )
+    #
+    #def reducer(x):
+    #    shape = K.shape(x)        
+    #    x = K.reshape(x, (-1, shape[-3] , shape[-2], shape[-1]) )      
+    #    y = K.depthwise_conv2d(x, K.constant(np.ones( (11,11,3,1) )), strides=(1, 1), padding='valid' )
+    #    y_shape = K.shape(y)
+    #    return K.reshape(y, (shape[0], y_shape[1], y_shape[2], y_shape[3] ) )
+    #        
+    #mean0 = reducer(y_true)
+    #mean1 = reducer(y_pred)
+    #luminance = mean0 * mean1    
+    #cs = y_true * y_pred
+    #
+    #result = K.function([y_true, y_pred],[luminance, cs]) ([image, image2])
+    #
+    #print (result)    
+    #import code
+    #code.interact(local=dict(globals(), **locals()))
+
+
     main()

mainscripts/Trainer.py

@@ -12,7 +12,7 @@ from utils import image_utils
 import cv2
 import models
 
-def trainerThread (input_queue, output_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name, save_interval_min=10, debug=False, **in_options):
+def trainerThread (input_queue, output_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name, save_interval_min=15, debug=False, **in_options):
 
     while True:
         try: 
@@ -39,10 +39,11 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
                         **in_options)
             
             is_reached_goal = model.is_reached_epoch_goal()
-            
+            is_upd_save_time_after_train = False
             def model_save():
                 if not debug and not is_reached_goal:
                     model.save()
+                    is_upd_save_time_after_train = True
             
             def send_preview():
                 if not debug:                        
@@ -65,10 +66,14 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
                 print('Starting. Press "Enter" to stop training and save model.')
  
             last_save_time = time.time()
+            
             for i in itertools.count(0,1):
                 if not debug:
                     if not is_reached_goal:
                         loss_string = model.train_one_epoch()     
+                        if is_upd_save_time_after_train:
+                            #save resets plaidML programs, so upd last_save_time only after plaidML rebuild them
+                            last_save_time = time.time()
                         
                         print (loss_string, end='\r')
                         if model.get_target_epoch() != 0 and model.is_reached_epoch_goal():

models/ModelBase.py

@@ -327,9 +327,9 @@ class ModelBase(object):
                 
         self.epoch += 1
 
-        if epoch_time >= 10000:
+        if epoch_time >= 10:
             #............."Saving... 
-            loss_string = "Training [#{0:06d}][{1:.5s}s]".format ( self.epoch, '{:0.4f}'.format(epoch_time / 1000) )
+            loss_string = "Training [#{0:06d}][{1:.5s}s]".format ( self.epoch, '{:0.4f}'.format(epoch_time) )
         else:
             loss_string = "Training [#{0:06d}][{1:04d}ms]".format ( self.epoch, int(epoch_time*1000) )
         for (loss_name, loss_value) in losses:

models/Model_SAE/Model.py

@@ -30,7 +30,7 @@ class SAEModel(ModelBase):
             self.options['resolution'] = input_int("Resolution (64,128 ?:help skip:128) : ", default_resolution, [64,128], help_message="More resolution requires more VRAM.")
             self.options['face_type'] = input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower()            
             self.options['learn_mask'] = input_bool ("Learn mask? (y/n, ?:help skip:y) : ", True, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.")
-            self.options['archi'] = input_str ("AE architecture (df, liae, ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae'], help_message="DF keeps faces more natural, while LIAE can fix overly different face shapes.").lower()            
+            self.options['archi'] = input_str ("AE architecture (df, liae, vg ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae','vg'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes. 'vg' - currently testing.").lower()
         else:
             self.options['resolution'] = self.options.get('resolution', default_resolution)
             self.options['face_type'] = self.options.get('face_type', default_face_type)
@@ -48,9 +48,13 @@ class SAEModel(ModelBase):
             
         if is_first_run:
             self.options['lighter_encoder'] = input_bool ("Use lightweight encoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight encoder is 35% faster, requires less VRAM, but sacrificing overall quality.")
+            
+            if self.options['archi'] != 'vg':
                 self.options['multiscale_decoder'] = input_bool ("Use multiscale decoder? (y/n, ?:help skip:y) : ", True, help_message="Multiscale decoder helps to get better details.")
         else:
             self.options['lighter_encoder'] = self.options.get('lighter_encoder', False)
+            
+            if self.options['archi'] != 'vg':
                 self.options['multiscale_decoder'] = self.options.get('multiscale_decoder', True)
             
         default_face_style_power = 0.0        
@@ -74,17 +78,19 @@ class SAEModel(ModelBase):
     #override
     def onInitialize(self, **in_options):
         exec(nnlib.import_all(), locals(), globals())
+        SAEModel.initialize_nn_functions()
         
         self.set_vram_batch_requirements({1.5:4})
         
         resolution = self.options['resolution']
         ae_dims = self.options['ae_dims']
         ed_ch_dims = self.options['ed_ch_dims']
-        adapt_k_size = False
         bgr_shape = (resolution, resolution, 3)
         mask_shape = (resolution, resolution, 1)
 
-        self.ms_count = ms_count = 3 if self.options['multiscale_decoder'] else 1
+        self.ms_count = ms_count = 3 if (self.options['archi'] != 'vg' and self.options['multiscale_decoder']) else 1
+        
+        masked_training = True
         
         epoch_alpha = Input( (1,) )
         warped_src = Input(bgr_shape)
@@ -101,7 +107,7 @@ class SAEModel(ModelBase):
         target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
 
         if self.options['archi'] == 'liae':
-            self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, adapt_k_size, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
+            self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
             
             enc_output_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] 
             
@@ -143,8 +149,8 @@ class SAEModel(ModelBase):
                 pred_dst_dstm = self.decoderm(warped_dst_inter_code)
                 pred_src_dstm = self.decoderm(warped_src_dst_inter_code)
 
-        else:
+        elif self.options['archi'] == 'df':
-            self.encoder = modelify(SAEModel.DFEncFlow(resolution, adapt_k_size, self.options['lighter_encoder'], ae_dims=ae_dims, ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
+            self.encoder = modelify(SAEModel.DFEncFlow(resolution, self.options['lighter_encoder'], ae_dims=ae_dims, ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
 
             dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] 
             
@@ -169,6 +175,38 @@ class SAEModel(ModelBase):
             pred_dst_dst = self.decoder_dst(warped_dst_code)
             pred_src_dst = self.decoder_src(warped_dst_code)
             
+            if self.options['learn_mask']:
+                pred_src_srcm = self.decoder_srcm(warped_src_code)
+                pred_dst_dstm = self.decoder_dstm(warped_dst_code)
+                pred_src_dstm = self.decoder_srcm(warped_dst_code)
+                
+        elif self.options['archi'] == 'vg':
+            self.encoder = modelify(SAEModel.VGEncFlow(resolution, self.options['lighter_encoder'], ae_dims=ae_dims, ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
+
+            dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] 
+            
+            self.decoder_src = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs)
+            self.decoder_dst = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs)
+            
+            if self.options['learn_mask']:
+                self.decoder_srcm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
+                self.decoder_dstm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
+            
+            if not self.is_first_run():
+                self.encoder.load_weights      (self.get_strpath_storage_for_file(self.encoderH5))
+                self.decoder_src.load_weights  (self.get_strpath_storage_for_file(self.decoder_srcH5))
+                self.decoder_dst.load_weights  (self.get_strpath_storage_for_file(self.decoder_dstH5))
+                if self.options['learn_mask']:
+                    self.decoder_srcm.load_weights (self.get_strpath_storage_for_file(self.decoder_srcmH5))
+                    self.decoder_dstm.load_weights (self.get_strpath_storage_for_file(self.decoder_dstmH5))
+            
+            warped_src_code = self.encoder (warped_src)
+            warped_dst_code = self.encoder (warped_dst)
+            pred_src_src = self.decoder_src(warped_src_code)
+            pred_dst_dst = self.decoder_dst(warped_dst_code)
+            pred_src_dst = self.decoder_src(warped_dst_code)
+
+
             if self.options['learn_mask']:
                 pred_src_srcm = self.decoder_srcm(warped_src_code)
                 pred_dst_dstm = self.decoder_dstm(warped_dst_code)
@@ -198,6 +236,15 @@ class SAEModel(ModelBase):
         target_dst_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_sigm_ar[i]  for i in range(len(target_dst_sigm_ar))]
         target_dst_anti_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i]  for i in range(len(target_dst_sigm_ar))]
     
+        pred_src_src_masked_ar = [ pred_src_src_sigm_ar[i] * target_srcm_sigm_ar[i]  for i in range(len(pred_src_src_sigm_ar))]
+        pred_dst_dst_masked_ar = [ pred_dst_dst_sigm_ar[i] * target_dstm_sigm_ar[i]  for i in range(len(pred_dst_dst_sigm_ar))]
+        
+        target_src_masked_ar_opt = target_src_masked_ar if masked_training else target_src_sigm_ar
+        target_dst_masked_ar_opt = target_dst_masked_ar if masked_training else target_dst_sigm_ar
+        
+        pred_src_src_masked_ar_opt = pred_src_src_masked_ar if masked_training else pred_src_src_sigm_ar
+        pred_dst_dst_masked_ar_opt = pred_dst_dst_masked_ar if masked_training else pred_dst_dst_sigm_ar
+        
         psd_target_dst_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_sigm_ar[i]  for i in range(len(pred_src_dst_sigm_ar))]
         psd_target_dst_anti_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i]  for i in range(len(pred_src_dst_sigm_ar))]
         
@@ -215,9 +262,9 @@ class SAEModel(ModelBase):
                     src_dst_mask_loss_train_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights
             
             if not self.options['pixel_loss']:
-                src_loss_batch = sum([ (  100*K.square( dssim(max_value=2.0)( target_src_masked_ar[i],  pred_src_src_sigm_ar[i] * target_srcm_sigm_ar[i] ) )) for i in range(len(target_src_masked_ar)) ])
+                src_loss_batch = sum([ (  100*K.square( dssim(max_value=2.0)( target_src_masked_ar_opt[i],  pred_src_src_masked_ar_opt[i] ) )) for i in range(len(target_src_masked_ar_opt)) ])
             else:
-                src_loss_batch = sum([ K.mean ( 100*K.square( target_src_masked_ar[i] - pred_src_src_sigm_ar[i] * target_srcm_sigm_ar[i] ), axis=[1,2,3]) for i in range(len(target_src_masked_ar)) ])
+                src_loss_batch = sum([ K.mean ( 100*K.square( target_src_masked_ar_opt[i] - pred_src_src_masked_ar_opt[i] ), axis=[1,2,3]) for i in range(len(target_src_masked_ar_opt)) ])
 
             src_loss = K.mean(src_loss_batch)
 
@@ -235,9 +282,9 @@ class SAEModel(ModelBase):
                 src_loss += bg_loss
 
             if not self.options['pixel_loss']:
-                dst_loss_batch = sum([ (  100*K.square(dssim(max_value=2.0)( target_dst_masked_ar[i],  pred_dst_dst_sigm_ar[i] * target_dstm_sigm_ar[i] ) )) for i in range(len(target_dst_masked_ar)) ])
+                dst_loss_batch = sum([ (  100*K.square(dssim(max_value=2.0)( target_dst_masked_ar_opt[i],  pred_dst_dst_masked_ar_opt[i] ) )) for i in range(len(target_dst_masked_ar_opt)) ])
             else:
-                dst_loss_batch = sum([ K.mean ( 100*K.square( target_dst_masked_ar[i] - pred_dst_dst_sigm_ar[i] * target_dstm_sigm_ar[i] ), axis=[1,2,3]) for i in range(len(target_dst_masked_ar)) ])
+                dst_loss_batch = sum([ K.mean ( 100*K.square( target_dst_masked_ar_opt[i] - pred_dst_dst_masked_ar_opt[i] ), axis=[1,2,3]) for i in range(len(target_dst_masked_ar_opt)) ])
                 
             dst_loss = K.mean(dst_loss_batch)
 
@@ -392,26 +439,66 @@ class SAEModel(ModelBase):
                                **in_options)
                                
     @staticmethod
-    def LIAEEncFlow(resolution, adapt_k_size, light_enc, ed_ch_dims=42):
+    def initialize_nn_functions():
         exec (nnlib.import_all(), locals(), globals())
         
-        k_size = resolution // 16 + 1 if adapt_k_size else 5
+        class ResidualBlock(object):
-        strides = resolution // 32 if adapt_k_size else 2
+            def __init__(self, filters, kernel_size=3, padding='same', use_reflection_padding=False):
+                self.filters = filters
+                self.kernel_size = kernel_size
+                self.padding = padding #if not use_reflection_padding else 'valid'
+                self.use_reflection_padding = use_reflection_padding
+                
+            def __call__(self, inp):
+                var_x = LeakyReLU(alpha=0.2)(inp)
+                
+                #if self.use_reflection_padding:
+                #    #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x)
+                    
+                var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x)
+                var_x = LeakyReLU(alpha=0.2)(var_x)
+                
+                #if self.use_reflection_padding:
+                #    #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x)
+                    
+                var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x)
+                var_x = Scale(gamma_init=keras.initializers.Constant(value=0.1))(var_x)
+                var_x = Add()([var_x, inp])
+                var_x = LeakyReLU(alpha=0.2)(var_x)
+                return var_x
+        SAEModel.ResidualBlock = ResidualBlock
 
         def downscale (dim):
             def func(x):
-                return LeakyReLU(0.1)(Conv2D(dim, k_size, strides=strides, padding='same')(x))
+                return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))
             return func 
+        SAEModel.downscale = downscale
         
         def downscale_sep (dim):
             def func(x):
-                return LeakyReLU(0.1)(SeparableConv2D(dim, k_size, strides=strides, padding='same')(x))
+                return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x))
             return func 
+        SAEModel.downscale_sep = downscale_sep
         
         def upscale (dim):
             def func(x):
-                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
+                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x)))
             return func 
+        SAEModel.upscale = upscale
+        
+        def to_bgr (output_nc):
+            def func(x):
+                return Conv2D(output_nc, kernel_size=5, padding='same', activation='tanh', kernel_initializer=RandomNormal(0, 0.02))(x)
+            return func
+        SAEModel.to_bgr = to_bgr
+        
+            
+    @staticmethod
+    def LIAEEncFlow(resolution, light_enc, ed_ch_dims=42):
+        exec (nnlib.import_all(), locals(), globals())
+        upscale = SAEModel.upscale
+        downscale = SAEModel.downscale
+        downscale_sep = SAEModel.downscale_sep
     
         def func(input):
             ed_dims = K.int_shape(input)[-1]*ed_ch_dims
@@ -434,13 +521,9 @@ class SAEModel(ModelBase):
     @staticmethod
     def LIAEInterFlow(resolution, ae_dims=256):
         exec (nnlib.import_all(), locals(), globals())
+        upscale = SAEModel.upscale
         lowest_dense_res=resolution // 16
         
-        def upscale (dim):
-            def func(x):
-                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func 
-        
         def func(input):   
             x = input[0]
             x = Dense(ae_dims)(x)
@@ -453,17 +536,10 @@ class SAEModel(ModelBase):
     @staticmethod
     def LIAEDecFlow(output_nc,ed_ch_dims=21, multiscale_count=1):
         exec (nnlib.import_all(), locals(), globals())
+        upscale = SAEModel.upscale
+        to_bgr = SAEModel.to_bgr
         ed_dims = output_nc * ed_ch_dims
             
-        def upscale (dim):
-            def func(x):
-                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func
-            
-        def to_bgr ():
-            def func(x):
-                return Conv2D(output_nc, kernel_size=5, padding='same', activation='tanh')(x)
-            return func
         def func(input):   
             x = input[0]
             
@@ -471,45 +547,28 @@ class SAEModel(ModelBase):
             x1     = upscale(ed_dims*8)( x )       
             
             if multiscale_count >= 3:
-                outputs += [ to_bgr() ( x1 ) ]  
+                outputs += [ to_bgr(output_nc) ( x1 ) ]  
                 
             x2     = upscale(ed_dims*4)( x1 )    
             
             if multiscale_count >= 2:
-                outputs += [ to_bgr() ( x2 ) ]
+                outputs += [ to_bgr(output_nc) ( x2 ) ]
                 
             x3     = upscale(ed_dims*2)( x2 )
             
-            outputs += [ to_bgr() ( x3 ) ]
+            outputs += [ to_bgr(output_nc) ( x3 ) ]
         
             return outputs
         return func
     
     @staticmethod
-    def DFEncFlow(resolution, adapt_k_size, light_enc, ae_dims=512, ed_ch_dims=42):
+    def DFEncFlow(resolution, light_enc, ae_dims=512, ed_ch_dims=42):
         exec (nnlib.import_all(), locals(), globals())
-        k_size = resolution // 16 + 1 if adapt_k_size else 5
+        upscale = SAEModel.upscale
-        strides = resolution // 32 if adapt_k_size else 2
+        downscale = SAEModel.downscale
+        downscale_sep = SAEModel.downscale_sep
         lowest_dense_res = resolution // 16
 
-        def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=True, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None):
-            return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint )
-
-        def downscale (dim):
-            def func(x):
-                return LeakyReLU(0.1)(Conv2D(dim, k_size, strides=strides, padding='same')(x))
-            return func 
-            
-        def downscale_sep (dim):
-            def func(x):
-                return LeakyReLU(0.1)(SeparableConv2D(dim, k_size, strides=strides, padding='same')(x))
-            return func 
-            
-        def upscale (dim):
-            def func(x):
-                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func   
-    
         def func(input):     
             x = input
             
@@ -536,20 +595,10 @@ class SAEModel(ModelBase):
     @staticmethod
     def DFDecFlow(output_nc, ed_ch_dims=21, multiscale_count=1):
         exec (nnlib.import_all(), locals(), globals())
+        upscale = SAEModel.upscale
+        to_bgr = SAEModel.to_bgr
         ed_dims = output_nc * ed_ch_dims
 
-        def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=True, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None):
-            return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint )
-
-        def upscale (dim):
-            def func(x):
-                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func
-            
-        def to_bgr ():
-            def func(x):
-                return Conv2D(output_nc, kernel_size=5, padding='same', activation='tanh')(x)
-            return func
         def func(input):   
             x = input[0]
             
@@ -557,18 +606,95 @@ class SAEModel(ModelBase):
             x1     = upscale(ed_dims*8)( x )       
             
             if multiscale_count >= 3:
-                outputs += [ to_bgr() ( x1 ) ]  
+                outputs += [ to_bgr(output_nc) ( x1 ) ]  
                 
             x2     = upscale(ed_dims*4)( x1 )    
             
             if multiscale_count >= 2:
-                outputs += [ to_bgr() ( x2 ) ]
+                outputs += [ to_bgr(output_nc) ( x2 ) ]
                 
             x3     = upscale(ed_dims*2)( x2 )
             
-            outputs += [ to_bgr() ( x3 ) ]
+            outputs += [ to_bgr(output_nc) ( x3 ) ]
         
             return outputs            
+        return func
+        
+    
+            
+    @staticmethod
+    def VGEncFlow(resolution, light_enc, ae_dims=512, ed_ch_dims=42):
+        exec (nnlib.import_all(), locals(), globals())
+        upscale = SAEModel.upscale
+        downscale = SAEModel.downscale
+        downscale_sep = SAEModel.downscale_sep
+        ResidualBlock = SAEModel.ResidualBlock
+        lowest_dense_res = resolution // 16
+            
+        def func(input):
+            x = input
+            ed_dims = K.int_shape(input)[-1]*ed_ch_dims
+            while np.modf(ed_dims / 4)[0] != 0.0:
+                ed_dims -= 1
+            
+            in_conv_filters = ed_dims if resolution <= 128 else ed_dims + (resolution//128)*ed_ch_dims
+            
+            x = tmp_x = Conv2D (in_conv_filters, kernel_size=5, strides=2, padding='same') (x)
+
+            for _ in range ( 8 if light_enc else 16 ):
+                x = ResidualBlock(ed_dims)(x)
+                
+            x = Add()([x, tmp_x])
+
+            x = downscale(ed_dims)(x)
+            x = SubpixelUpscaler()(x)
+            
+            x = downscale(ed_dims)(x)
+            x = SubpixelUpscaler()(x)
+            
+            x = downscale(ed_dims)(x)           
+            if light_enc:
+                x = downscale_sep (ed_dims*2)(x)
+            else:
+                x = downscale (ed_dims*2)(x)
+                
+            x = downscale(ed_dims*4)(x)
+            
+            if light_enc:
+                x = downscale_sep (ed_dims*8)(x)
+            else:
+                x = downscale (ed_dims*8)(x)
+            
+            x = Dense(ae_dims)(Flatten()(x))
+            x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
+            x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
+            x = upscale(ae_dims)(x)
+            return x
+            
+        return func
+        
+    @staticmethod
+    def VGDecFlow(output_nc, ed_ch_dims=21, multiscale_count=1):
+        exec (nnlib.import_all(), locals(), globals())
+        upscale = SAEModel.upscale        
+        to_bgr = SAEModel.to_bgr
+        ResidualBlock = SAEModel.ResidualBlock
+        ed_dims = output_nc * ed_ch_dims
+  
+        def func(input):
+            x = input[0]
+            
+            x = upscale( ed_dims*8 )(x)
+            x = ResidualBlock( ed_dims*8 )(x)
+            
+            x = upscale( ed_dims*4 )(x)
+            x = ResidualBlock( ed_dims*4 )(x)
+            
+            x = upscale( ed_dims*2 )(x)
+            x = ResidualBlock( ed_dims*2 )(x)
+        
+            x = to_bgr(output_nc) (x)        
+            return x
             
         return func
         

nnlib/nnlib.py

@@ -62,6 +62,7 @@ Lambda = keras.layers.Lambda
 Add = keras.layers.Add
 Concatenate = keras.layers.Concatenate
 
+
 Flatten = keras.layers.Flatten
 Reshape = keras.layers.Reshape
 
@@ -77,9 +78,11 @@ gaussian_blur = nnlib.gaussian_blur
 style_loss = nnlib.style_loss
 dssim = nnlib.dssim
 
-#ReflectionPadding2D = nnlib.ReflectionPadding2D
+
 PixelShuffler = nnlib.PixelShuffler
 SubpixelUpscaler = nnlib.SubpixelUpscaler
+Scale = nnlib.Scale
+#ReflectionPadding2D = nnlib.ReflectionPadding2D
 #AddUniformNoise = nnlib.AddUniformNoise
 """
     code_import_keras_contrib_string = \
@@ -184,8 +187,9 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
         if 'TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1':        
             suppressor.__exit__()
             
-        nnlib.__initialize_keras_functions()  
         nnlib.code_import_keras = compile (nnlib.code_import_keras_string,'','exec')
+        nnlib.__initialize_keras_functions()  
+        
         return nnlib.code_import_keras
         
     @staticmethod
@@ -395,8 +399,40 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
         
         nnlib.PixelShuffler = PixelShuffler
         nnlib.SubpixelUpscaler = PixelShuffler
-        '''
 
+        class Scale(keras.layers.Layer):
+            """
+            GAN Custom Scal Layer
+            Code borrows from https://github.com/flyyufelix/cnn_finetune
+            """
+            def __init__(self, weights=None, axis=-1, gamma_init='zero', **kwargs):
+                self.axis = axis
+                self.gamma_init = keras.initializers.get(gamma_init)
+                self.initial_weights = weights
+                super(Scale, self).__init__(**kwargs)
+
+            def build(self, input_shape):
+                self.input_spec = [keras.engine.InputSpec(shape=input_shape)]
+
+                # Compatibility with TensorFlow >= 1.0.0
+                self.gamma = K.variable(self.gamma_init((1,)), name='{}_gamma'.format(self.name))
+                self.trainable_weights = [self.gamma]
+
+                if self.initial_weights is not None:
+                    self.set_weights(self.initial_weights)
+                    del self.initial_weights
+
+            def call(self, x, mask=None):
+                return self.gamma * x
+
+            def get_config(self):
+                config = {"axis": self.axis}
+                base_config = super(Scale, self).get_config()
+                return dict(list(base_config.items()) + list(config.items()))
+        nnlib.Scale = Scale
+        
+        '''        
+        not implemented in plaidML
         class ReflectionPadding2D(keras.layers.Layer):
             def __init__(self, padding=(1, 1), **kwargs):
                 self.padding = tuple(padding)
@@ -411,27 +447,10 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
                 w_pad,h_pad = self.padding
                 return tf.pad(x, [[0,0], [h_pad,h_pad], [w_pad,w_pad], [0,0] ], 'REFLECT')
         nnlib.ReflectionPadding2D = ReflectionPadding2D     
-
-       
-        class AddUniformNoise(keras.layers.Layer):
-            def __init__(self, power=1.0, minval=-1.0, maxval=1.0, **kwargs):
-                super(AddUniformNoise, self).__init__(**kwargs)
-                self.power = power
-                self.supports_masking = True
-                self.minval = minval
-                self.maxval = maxval
-
-            def call(self, inputs, training=None):
-                def noised():
-                    return inputs + self.power*K.random_uniform(shape=K.shape(inputs), minval=self.minval, maxval=self.maxval)
-                return K.in_train_phase(noised, inputs, training=training)
-
-            def get_config(self):
-                config = {'power': self.power, 'minval': self.minval, 'maxval': self.maxval}
-                base_config = super(AddUniformNoise, self).get_config()
-                return dict(list(base_config.items()) + list(config.items()))
-        nnlib.AddUniformNoise = AddUniformNoise       
         '''   
+        
+        
+             
     @staticmethod
     def import_keras_contrib(device_config = None):
         if nnlib.keras_contrib is not None:
@@ -489,6 +508,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
                     return 10*dssim() (y_true*mask, y_pred*mask)                    
         nnlib.DSSIMMSEMaskLoss = DSSIMMSEMaskLoss
         
+
         '''
         def ResNet(output_nc, use_batch_norm, ngf=64, n_blocks=6, use_dropout=False):
             exec (nnlib.import_all(), locals(), globals())