added preview_samples, learning_rate and cpu cores as option

models/Model_AMP/Model.py

@@ -59,6 +59,8 @@ class AMPModel(ModelBase):
         default_clipgrad           = self.options['clipgrad']           = self.load_or_def_option('clipgrad', False)
         default_usefp16            = self.options['use_fp16']           = self.load_or_def_option('use_fp16', False)
         default_cpu_cap            = self.options['cpu_cap']            = self.load_or_def_option('default_cpu_cap', 8)
+        default_preview_samples    = self.options['preview_samples']    = self.load_or_def_option('preview_samples', 4)
+        default_lr_modifier        = self.options['lr_modifier']        = self.load_or_def_option('lr_modifier', 0)
 
         ask_override = False if self.read_from_conf else self.ask_override()
         if self.is_first_run() or ask_override:
@@ -67,6 +69,7 @@ class AMPModel(ModelBase):
                 self.ask_session_name()
                 self.ask_maximum_n_backups()
                 self.ask_write_preview_history()
+                self.options['preview_samples'] = np.clip ( io.input_int ("Number of samples to preview", default_cpu_cap, add_info="1 - 16", help_message="Typical fine value is 4"), 1, 16 )
                 self.ask_target_iter()
                 self.ask_retraining_samples()
                 self.ask_random_src_flip()
@@ -125,6 +128,8 @@ class AMPModel(ModelBase):
                     self.options['blur_out_mask'] = io.input_bool ("Blur out mask", default_blur_out_mask, help_message='Blurs nearby area outside of applied face mask of training samples. The result is the background near the face is smoothed and less noticeable on swapped face. The exact xseg mask in src and dst faceset is required.')
                 
                 self.options['loss_function'] = io.input_str(f"Loss function", default_loss_function, ['SSIM', 'MS-SSIM', 'MS-SSIM+L1'], help_message="Change loss function used for image quality assessment.")
+                self.options['lr_modifier'] = np.clip (io.input_int("Learningrate factor", default_lr_modifier, add_info="-100 .. 100", help_message="Modify the Learning rate: 100 == multipy by 4, -100 == divide by 4"), -100, 100)
+
                 self.options['lr_dropout']  = io.input_str (f"Use learning rate dropout", default_lr_dropout, ['n','y','cpu'], help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations. Enabled it before `disable random warp` and before GAN. \nn - disabled.\ny - enabled\ncpu - enabled on CPU. This allows not to use extra VRAM, sacrificing 20% time of iteration.")
 
         default_gan_power          = self.options['gan_power']          = self.load_or_def_option('gan_power', 0.0)
@@ -146,6 +151,7 @@ class AMPModel(ModelBase):
                 self.options['random_blur'] = io.input_bool("Enable random blur of samples", default_random_blur, help_message="")
                 self.options['random_jpeg'] = io.input_bool("Enable random jpeg compression of samples", default_random_jpeg, help_message="")
                 
+                
                 #self.options['random_hsv_power'] = np.clip ( io.input_number ("Random hue/saturation/light intensity", default_random_hsv_power, add_info="0.0 .. 0.3", help_message="Random hue/saturation/light intensity applied to the src face set only at the input of the neural network. Stabilizes color perturbations during face swapping. Reduces the quality of the color transfer by selecting the closest one in the src faceset. Thus the src faceset must be diverse enough. Typical fine value is 0.05"), 0.0, 0.3 )
 
                 self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 5.0", help_message="Forces the neural network to learn small details of the face. Enable it only when the face is trained enough with random_warp(off), and don't disable. The higher the value, the higher the chances of artifacts. Typical fine value is 0.1"), 0.0, 5.0 )
@@ -388,6 +394,13 @@ class AMPModel(ModelBase):
                         self.model_filename_list += [ [self.GAN, 'GAN.npy'] ]
             
                 # Initialize optimizers
+                lr_modifier = self.options['lr_modifier']
+                if lr_modifier == 0:
+                    lr = 5e-5 
+                elif lr_modifier < 0:
+                    lr = 5e-5 / abs( lr_modifier * 4/100 ) 
+                else:
+                    lr = 5e-5 * abs( lr_modifier * 4/100 ) 
                 clipnorm = 1.0 if self.options['clipgrad'] else 0.0
                 if self.options['lr_dropout'] in ['y','cpu']:
                     lr_cos = 500
@@ -398,17 +411,17 @@ class AMPModel(ModelBase):
                 self.G_weights = self.encoder.get_weights() + self.decoder.get_weights()
 
                 OptimizerClass = nn.AdaBelief if adabelief else nn.RMSprop
-                self.src_dst_opt = OptimizerClass(lr=5e-5, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
+                self.src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
                 self.src_dst_opt.initialize_variables (self.G_weights, vars_on_cpu=optimizer_vars_on_cpu)
                 self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
 
                 if gan_power != 0:
                     if self.options['gan_version'] == 2:
-                        self.GAN_opt = OptimizerClass(lr=5e-5, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='D_src_dst_opt')
+                        self.GAN_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='D_src_dst_opt')
                         self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
                         self.model_filename_list += [ (self.GAN_opt, 'D_src_v2_opt.npy') ]
                     else:
-                        self.GAN_opt = OptimizerClass(lr=5e-5, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt')
+                        self.GAN_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt')
                         self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
                         self.model_filename_list += [ (self.GAN_opt, 'GAN_opt.npy') ]
 
@@ -920,26 +933,31 @@ class AMPModel(ModelBase):
 
         target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]
 
-        n_samples = min(4, self.get_batch_size(), 800 // self.resolution )
+        #n_samples = min(4, self.get_batch_size(), 800 // self.resolution )
+        n_samples = min(self.get_batch_size(), self.options['preview_samples'])
 
         result = []
 
         i = np.random.randint(n_samples) if not for_history else 0
 
+        for i in range(n_samples if not for_history else 1):
             if filenames is not None and len(filenames) > 0:
                 S[i] = label_face_filename(S[i], filenames[0][i])
                 D[i] = label_face_filename(D[i], filenames[1][i])
-
+        st = []        
-        st =  [ np.concatenate ((S[i],  D[i],  DD[i]*DDM_000[i]), axis=1) ]
+        for i in range(n_samples):
+            st +=  [ np.concatenate ((S[i],  D[i],  DD[i]*DDM_000[i]), axis=1) ]
             st += [ np.concatenate ((SS[i], DD[i], SD_100[i] ), axis=1) ]
 
         result += [ ('AMP morph 1.0', np.concatenate (st, axis=0 )), ]
-
+        st = []    
-        st =  [ np.concatenate ((DD[i], SD_025[i],  SD_050[i]), axis=1) ]
+        for i in range(n_samples):
+            st +=  [ np.concatenate ((DD[i], SD_025[i],  SD_050[i]), axis=1) ]
             st += [ np.concatenate ((SD_065[i], SD_075[i], SD_100[i]), axis=1) ]
         result += [ ('AMP morph list', np.concatenate (st, axis=0 )), ]
-
+        st = []    
-        st =  [ np.concatenate ((DD[i], SD_025[i]*DDM_025[i]*SDM_025[i],  SD_050[i]*DDM_050[i]*SDM_050[i]), axis=1) ]
+        for i in range(n_samples):
+            st += [ np.concatenate ((DD[i], SD_025[i]*DDM_025[i]*SDM_025[i],  SD_050[i]*DDM_050[i]*SDM_050[i]), axis=1) ]
             st += [ np.concatenate ((SD_065[i]*DDM_065[i]*SDM_065[i], SD_075[i]*DDM_075[i]*SDM_075[i], SD_100[i]*DDM_100[i]*SDM_100[i]), axis=1) ]
         result += [ ('AMP morph list masked', np.concatenate (st, axis=0 )), ]
 

models/Model_AMP/config_schema.json

@@ -102,6 +102,11 @@
             "MS-SSIM+L1"
           ]
         },
+        "lr_modifier": {
+          "type": "integer",
+          "minimum": -100,
+          "maximum": 100
+        },
         "random_warp": {
           "type": "boolean"
         },
@@ -234,6 +239,7 @@
         "gan_version",
         "loss_function",
         "lr_dropout",
+        "lr_modifier",
         "masked_training",
         "maximum_n_backups",
         "models_opt_on_gpu",

models/Model_SAEHD/Model.py

@@ -73,7 +73,7 @@ class SAEHDModel(ModelBase):
         default_clipgrad           = self.options['clipgrad']           = self.load_or_def_option('clipgrad', False)
         default_pretrain           = self.options['pretrain']           = self.load_or_def_option('pretrain', False)
         default_cpu_cap            = self.options['cpu_cap']            = self.load_or_def_option('cpu_cap', 8)
-        default_preview_samples    = self.options['default_preview_samples'] = self.load_or_def_option('preview_samples', 4)
+        default_preview_samples    = self.options['preview_samples']    = self.load_or_def_option('preview_samples', 4)
         default_lr_modifier        = self.options['lr_modifier']        = self.load_or_def_option('lr_modifier', 0)
 
         ask_override = False if self.read_from_conf else self.ask_override()

models/Model_SAEHD/config_schema.json

@@ -95,6 +95,11 @@
             "MS-SSIM+L1"
           ]
         },
+        "lr_modifier": {
+          "type": "integer",
+          "minimum": -100,
+          "maximum": 100
+        },
         "random_warp": {
           "type": "boolean"
         },
@@ -155,6 +160,11 @@
         "pretrain": {
           "type": "boolean"
         },
+        "preview_samples": {
+          "type": "integer",
+          "minimum": 1,
+          "maximum": 16
+        },
         "session_name": {
           "type": "string"
         },
@@ -243,11 +253,13 @@
         "gan_version",
         "loss_function",
         "lr_dropout",
+        "lr_modifier",
         "masked_training",
         "maximum_n_backups",
         "models_opt_on_gpu",
         "mouth_prio",
         "pretrain",
+        "preview_samples",
         "random_blur",
         "random_color",
         "random_downsample",