From 6d89d7fa4c1794279993d68729937c3fa773f89c Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 26 Jun 2021 10:42:11 +0400 Subject: [PATCH 01/30] _ --- .vscode/launch.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index f8857c1..6dc974d 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -12,7 +12,7 @@ "type": "python", "request": "launch", "program": "${env:DFL_ROOT}\\main.py", - "pythonPath": "${env:PYTHONEXECUTABLE}", + "python": "${env:PYTHONEXECUTABLE}", "cwd": "${env:WORKSPACE}", "console": "integratedTerminal", "args": ["train", From 578319184932fca4684d183b5f585b63b82926ea Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 26 Jun 2021 10:44:41 +0400 Subject: [PATCH 02/30] AMP: code refactoring, fix preview history added dumpdflive command --- main.py | 13 ++- mainscripts/Trainer.py | 13 +-- models/ModelBase.py | 12 +- models/Model_AMP/Model.py | 210 ++++++++++++++++------------------ models/Model_Quick96/Model.py | 2 +- models/Model_SAEHD/Model.py | 31 +++-- models/Model_XSeg/Model.py | 2 +- requirements-colab.txt | 3 +- requirements-cuda.txt | 1 + 9 files changed, 143 insertions(+), 144 deletions(-) diff --git a/main.py b/main.py index d9fd4fd..04337af 100644 --- a/main.py +++ b/main.py @@ -127,7 +127,6 @@ if __name__ == "__main__": 'silent_start' : arguments.silent_start, 'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ], 'debug' : arguments.debug, - 'dump_ckpt' : arguments.dump_ckpt, } from mainscripts import Trainer Trainer.main(**kwargs) @@ -145,11 +144,19 @@ if __name__ == "__main__": p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.") p.add_argument('--silent-start', action="store_true", dest="silent_start", default=False, help="Silent start. Automatically chooses Best GPU and last used model.") - p.add_argument('--dump-ckpt', action="store_true", dest="dump_ckpt", default=False, help="Dump the model to ckpt format.") - p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+') p.set_defaults (func=process_train) + + def process_dumpdflive(arguments): + osex.set_process_lowest_prio() + from mainscripts import DumpDFLive + DumpDFLive.main(model_class_name = arguments.model_name, saved_models_path = Path(arguments.model_dir)) + + p = subparsers.add_parser( "dumpdflive", help="Dump model to use in DFLive.") + p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.") + p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.") + p.set_defaults (func=process_dumpdflive) def process_merge(arguments): osex.set_process_lowest_prio() diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py index 780cc48..5cc8acf 100644 --- a/mainscripts/Trainer.py +++ b/mainscripts/Trainer.py @@ -27,7 +27,6 @@ def trainerThread (s2c, c2s, e, silent_start=False, execute_programs = None, debug=False, - dump_ckpt=False, **kwargs): while True: try: @@ -43,12 +42,9 @@ def trainerThread (s2c, c2s, e, if not saved_models_path.exists(): saved_models_path.mkdir(exist_ok=True, parents=True) - - if dump_ckpt: - cpu_only=True - + model = models.import_model(model_class_name)( - is_training=not dump_ckpt, + is_training=True, saved_models_path=saved_models_path, training_data_src_path=training_data_src_path, training_data_dst_path=training_data_dst_path, @@ -61,11 +57,6 @@ def trainerThread (s2c, c2s, e, silent_start=silent_start, debug=debug) - if dump_ckpt: - e.set() - model.dump_ckpt() - break - is_reached_goal = model.is_reached_iter_goal() shared_state = { 'after_save' : False } diff --git a/models/ModelBase.py b/models/ModelBase.py index ce219ab..598c19a 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -232,7 +232,7 @@ class ModelBase(object): preview_id_counter = 0 while not choosed: self.sample_for_preview = self.generate_next_samples() - previews = self.get_static_previews() + previews = self.get_history_previews() io.show_image( wnd_name, ( previews[preview_id_counter % len(previews) ][1] *255).astype(np.uint8) ) @@ -258,7 +258,7 @@ class ModelBase(object): self.sample_for_preview = self.generate_next_samples() try: - self.get_static_previews() + self.get_history_previews() except: self.sample_for_preview = self.generate_next_samples() @@ -347,7 +347,7 @@ class ModelBase(object): return ( ('loss_src', 0), ('loss_dst', 0) ) #overridable - def onGetPreview(self, sample): + def onGetPreview(self, sample, for_history=False): #you can return multiple previews #return [ ('preview_name',preview_rgb), ... ] return [] @@ -377,8 +377,8 @@ class ModelBase(object): def get_previews(self): return self.onGetPreview ( self.last_sample ) - def get_static_previews(self): - return self.onGetPreview (self.sample_for_preview) + def get_history_previews(self): + return self.onGetPreview (self.sample_for_preview, for_history=True) def get_preview_history_writer(self): if self.preview_history_writer is None: @@ -484,7 +484,7 @@ class ModelBase(object): plist += [ (bgr, self.get_strpath_storage_for_file('preview_%s.jpg' % (name) ) ) ] if self.write_preview_history: - previews = self.get_static_previews() + previews = self.get_history_previews() for i in range(len(previews)): name, bgr = previews[i] path = self.preview_history_path / name diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 9347944..5b75036 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -87,7 +87,7 @@ class AMPModel(ModelBase): d_mask_dims = np.clip ( io.input_int("Decoder mask dimensions", default_d_mask_dims, add_info="16-256", help_message="Typical mask dimensions = decoder dimensions / 3. If you manually cut out obstacles from the dst mask, you can increase this parameter to achieve better quality." ), 16, 256 ) self.options['d_mask_dims'] = d_mask_dims + d_mask_dims % 2 - + morph_factor = np.clip ( io.input_number ("Morph factor.", default_morph_factor, add_info="0.1 .. 0.5", help_message="The smaller the value, the more src-like facial expressions will appear. The larger the value, the less space there is to train a large dst faceset in the neural network. Typical fine value is 0.33"), 0.1, 0.5 ) self.options['morph_factor'] = morph_factor @@ -121,9 +121,9 @@ class AMPModel(ModelBase): self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") self.options['clipgrad'] = io.input_bool ("Enable gradient clipping", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.") - + self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly. Forces random_warp=N, random_flips=Y, gan_power=0.0, lr_dropout=N, uniform_yaw=Y") - + self.gan_model_changed = (default_gan_patch_size != self.options['gan_patch_size']) or (default_gan_dims != self.options['gan_dims']) self.pretrain_just_disabled = (default_pretrain == True and self.options['pretrain'] == False) @@ -137,34 +137,26 @@ class AMPModel(ModelBase): self.resolution = resolution = self.options['resolution'] + input_ch=3 + ae_dims = self.ae_dims = self.options['ae_dims'] + e_dims = self.options['e_dims'] + d_dims = self.options['d_dims'] + d_mask_dims = self.options['d_mask_dims'] lowest_dense_res = self.lowest_dense_res = resolution // 32 class Downscale(nn.ModelBase): - def __init__(self, in_ch, out_ch, kernel_size=5, *kwargs ): - self.in_ch = in_ch - self.out_ch = out_ch - self.kernel_size = kernel_size - super().__init__(*kwargs) - - def on_build(self, *args, **kwargs ): - self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME') + def on_build(self, in_ch, out_ch, kernel_size=5 ): + self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME') def forward(self, x): - x = self.conv1(x) - x = tf.nn.leaky_relu(x, 0.1) - return x - - def get_out_ch(self): - return self.out_ch + return tf.nn.leaky_relu(self.conv1(x), 0.1) class Upscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME') + self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME') def forward(self, x): - x = self.conv1(x) - x = tf.nn.leaky_relu(x, 0.1) - x = nn.depth_to_space(x, 2) + x = nn.depth_to_space(tf.nn.leaky_relu(self.conv1(x), 0.1), 2) return x class ResidualBlock(nn.ModelBase): @@ -180,15 +172,15 @@ class AMPModel(ModelBase): return x class Encoder(nn.ModelBase): - def on_build(self, in_ch, e_ch, ae_ch): - self.down1 = Downscale(in_ch, e_ch, kernel_size=5) - self.res1 = ResidualBlock(e_ch) - self.down2 = Downscale(e_ch, e_ch*2, kernel_size=5) - self.down3 = Downscale(e_ch*2, e_ch*4, kernel_size=5) - self.down4 = Downscale(e_ch*4, e_ch*8, kernel_size=5) - self.down5 = Downscale(e_ch*8, e_ch*8, kernel_size=5) - self.res5 = ResidualBlock(e_ch*8) - self.dense1 = nn.Dense( lowest_dense_res*lowest_dense_res*e_ch*8, ae_ch ) + def on_build(self): + self.down1 = Downscale(input_ch, e_dims, kernel_size=5) + self.res1 = ResidualBlock(e_dims) + self.down2 = Downscale(e_dims, e_dims*2, kernel_size=5) + self.down3 = Downscale(e_dims*2, e_dims*4, kernel_size=5) + self.down4 = Downscale(e_dims*4, e_dims*8, kernel_size=5) + self.down5 = Downscale(e_dims*8, e_dims*8, kernel_size=5) + self.res5 = ResidualBlock(e_dims*8) + self.dense1 = nn.Dense( lowest_dense_res*lowest_dense_res*e_dims*8, ae_dims ) def forward(self, inp): x = inp @@ -199,53 +191,45 @@ class AMPModel(ModelBase): x = self.down4(x) x = self.down5(x) x = self.res5(x) - x = nn.flatten(x) - x = nn.pixel_norm(x, axes=-1) + x = nn.pixel_norm(nn.flatten(x), axes=-1) x = self.dense1(x) return x class Inter(nn.ModelBase): - def __init__(self, ae_ch, ae_out_ch, **kwargs): - self.ae_ch, self.ae_out_ch = ae_ch, ae_out_ch - super().__init__(**kwargs) - def on_build(self): - ae_ch, ae_out_ch = self.ae_ch, self.ae_out_ch - self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch ) + self.dense2 = nn.Dense( ae_dims, lowest_dense_res * lowest_dense_res * ae_dims ) def forward(self, inp): x = inp x = self.dense2(x) - x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch) + x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, ae_dims) return x - def get_out_ch(self): - return self.ae_out_ch class Decoder(nn.ModelBase): - def on_build(self, in_ch, d_ch, d_mask_ch ): - self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3) - self.upscale1 = Upscale(d_ch*8, d_ch*8, kernel_size=3) - self.upscale2 = Upscale(d_ch*8, d_ch*4, kernel_size=3) - self.upscale3 = Upscale(d_ch*4, d_ch*2, kernel_size=3) + def on_build(self ): + self.upscale0 = Upscale(ae_dims, d_dims*8, kernel_size=3) + self.upscale1 = Upscale(d_dims*8, d_dims*8, kernel_size=3) + self.upscale2 = Upscale(d_dims*8, d_dims*4, kernel_size=3) + self.upscale3 = Upscale(d_dims*4, d_dims*2, kernel_size=3) - self.res0 = ResidualBlock(d_ch*8, kernel_size=3) - self.res1 = ResidualBlock(d_ch*8, kernel_size=3) - self.res2 = ResidualBlock(d_ch*4, kernel_size=3) - self.res3 = ResidualBlock(d_ch*2, kernel_size=3) + self.res0 = ResidualBlock(d_dims*8, kernel_size=3) + self.res1 = ResidualBlock(d_dims*8, kernel_size=3) + self.res2 = ResidualBlock(d_dims*4, kernel_size=3) + self.res3 = ResidualBlock(d_dims*2, kernel_size=3) - self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3) - self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*8, kernel_size=3) - self.upscalem2 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3) - self.upscalem3 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3) - self.upscalem4 = Upscale(d_mask_ch*2, d_mask_ch*1, kernel_size=3) - self.out_convm = nn.Conv2D( d_mask_ch*1, 1, kernel_size=1, padding='SAME') + self.upscalem0 = Upscale(ae_dims, d_mask_dims*8, kernel_size=3) + self.upscalem1 = Upscale(d_mask_dims*8, d_mask_dims*8, kernel_size=3) + self.upscalem2 = Upscale(d_mask_dims*8, d_mask_dims*4, kernel_size=3) + self.upscalem3 = Upscale(d_mask_dims*4, d_mask_dims*2, kernel_size=3) + self.upscalem4 = Upscale(d_mask_dims*2, d_mask_dims*1, kernel_size=3) + self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME') - self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME') - self.out_conv1 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME') - self.out_conv2 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME') - self.out_conv3 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME') + self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME') + self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME') + self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME') + self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME') def forward(self, inp): z = inp @@ -280,27 +264,24 @@ class AMPModel(ModelBase): eyes_mouth_prio = self.options['eyes_mouth_prio'] - ae_dims = self.ae_dims = self.options['ae_dims'] - e_dims = self.options['e_dims'] - d_dims = self.options['d_dims'] - d_mask_dims = self.options['d_mask_dims'] + morph_factor = self.options['morph_factor'] - + pretrain = self.pretrain = self.options['pretrain'] if self.pretrain_just_disabled: self.set_iter(0) - + self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power'] random_warp = False if self.pretrain else self.options['random_warp'] random_src_flip = self.random_src_flip if not self.pretrain else True random_dst_flip = self.random_dst_flip if not self.pretrain else True - + if self.pretrain: self.options_show_override['gan_power'] = 0.0 self.options_show_override['random_warp'] = False self.options_show_override['lr_dropout'] = 'n' self.options_show_override['uniform_yaw'] = True - + masked_training = self.options['masked_training'] ct_mode = self.options['ct_mode'] if ct_mode == 'none': @@ -310,7 +291,7 @@ class AMPModel(ModelBase): models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0' optimizer_vars_on_cpu = models_opt_device=='/CPU:0' - input_ch=3 + bgr_shape = self.bgr_shape = nn.get4Dshape(resolution,resolution,input_ch) mask_shape = nn.get4Dshape(resolution,resolution,1) self.model_filename_list = [] @@ -333,10 +314,10 @@ class AMPModel(ModelBase): # Initializing model classes with tf.device (models_opt_device): - self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, ae_ch=ae_dims, name='encoder') - self.inter_src = Inter(ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter_src') - self.inter_dst = Inter(ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter_dst') - self.decoder = Decoder(in_ch=ae_dims, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder') + self.encoder = Encoder(name='encoder') + self.inter_src = Inter(name='inter_src') + self.inter_dst = Inter(name='inter_dst') + self.decoder = Decoder(name='decoder') self.model_filename_list += [ [self.encoder, 'encoder.npy'], [self.inter_src, 'inter_src.npy'], @@ -351,7 +332,7 @@ class AMPModel(ModelBase): # Initialize optimizers lr=5e-5 lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] and not self.pretrain else 1.0 - + clipnorm = 1.0 if self.options['clipgrad'] else 0.0 self.all_weights = self.encoder.get_weights() + self.inter_src.get_weights() + self.inter_dst.get_weights() + self.decoder.get_weights() @@ -386,8 +367,6 @@ class AMPModel(ModelBase): gpu_src_losses = [] gpu_dst_losses = [] gpu_G_loss_gvs = [] - gpu_GAN_loss_gvs = [] - gpu_D_code_loss_gvs = [] gpu_D_src_dst_loss_gvs = [] for gpu_id in range(gpu_count): @@ -407,7 +386,7 @@ class AMPModel(ModelBase): # process model tensors gpu_src_code = self.encoder (gpu_warped_src) gpu_dst_code = self.encoder (gpu_warped_dst) - + if pretrain: gpu_src_inter_src_code = self.inter_src (gpu_src_code) gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) @@ -454,7 +433,7 @@ class AMPModel(ModelBase): gpu_pred_src_src_anti_masked = gpu_pred_src_src*(1.0-gpu_target_srcm_blur) gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst gpu_pred_dst_dst_anti_masked = gpu_pred_dst_dst*(1.0-gpu_target_dstm_blur) - + if resolution < 256: gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) else: @@ -481,12 +460,12 @@ class AMPModel(ModelBase): gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) else: gpu_src_loss = gpu_dst_loss - + gpu_src_losses += [gpu_src_loss] - + if pretrain: gpu_G_loss = gpu_dst_loss - else: + else: gpu_G_loss = gpu_src_loss + gpu_dst_loss def DLossOnes(logits): @@ -537,8 +516,6 @@ class AMPModel(ModelBase): if gan_power != 0: src_D_src_dst_loss_gv_op = self.GAN_opt.get_update_op (nn.average_gv_list(gpu_D_src_dst_loss_gvs) ) - #GAN_loss_gv_op = self.src_dst_opt.get_update_op (nn.average_gv_list(gpu_GAN_loss_gvs) ) - # Initializing training and view functions def src_dst_train(warped_src, target_src, target_srcm, target_srcm_em, \ @@ -569,7 +546,6 @@ class AMPModel(ModelBase): self.target_dstm_em:target_dstm_em}) self.D_src_dst_train = D_src_dst_train - def AE_view(warped_src, warped_dst, morph_value): return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm], feed_dict={self.warped_src:warped_src, self.warped_dst:warped_dst, self.morph_value_t:[morph_value] }) @@ -605,13 +581,11 @@ class AMPModel(ModelBase): if self.is_training and gan_power != 0 and model == self.GAN: if self.gan_model_changed: do_init = True - + if not do_init: do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) ) if do_init: model.init_weights() - - ############### # initializing sample generators @@ -621,7 +595,6 @@ class AMPModel(ModelBase): random_ct_samples_path=training_data_dst_path if ct_mode is not None and not self.pretrain else None - cpu_count = min(multiprocessing.cpu_count(), 8) src_generators_count = cpu_count // 2 dst_generators_count = cpu_count // 2 @@ -654,9 +627,13 @@ class AMPModel(ModelBase): self.last_dst_samples_loss = [] if self.pretrain_just_disabled: self.update_sample_for_preview(force_new=True) - - def dump_ckpt(self): + + def dump_dflive (self): + output_path=self.get_strpath_storage_for_file('model.dflive') + + io.log_info(f'Dumping .dflive to {output_path}') + tf = nn.tf with tf.device (nn.tf_default_device_name): warped_dst = tf.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face') @@ -667,9 +644,9 @@ class AMPModel(ModelBase): gpu_dst_inter_src_code = self.inter_src ( gpu_dst_code) gpu_dst_inter_dst_code = self.inter_dst ( gpu_dst_code) - ae_dims_slice = tf.cast(self.ae_dims*morph_value[0], tf.int32) - gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, ae_dims_slice , self.lowest_dense_res, self.lowest_dense_res]), - tf.slice(gpu_dst_inter_dst_code, [0,ae_dims_slice,0,0], [-1,self.ae_dims-ae_dims_slice, self.lowest_dense_res,self.lowest_dense_res]) ), 1 ) + inter_dims_slice = tf.cast(self.inter_dims*morph_value[0], tf.int32) + gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , self.inter_res, self.inter_res]), + tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,self.inter_dims-inter_dims_slice, self.inter_res,self.inter_res]) ), 1 ) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code) _, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code) @@ -688,9 +665,15 @@ class AMPModel(ModelBase): ['out_face_mask','out_celeb_face','out_celeb_face_mask'] ) - pb_filepath = self.get_strpath_storage_for_file('.pb') - with tf.gfile.GFile(pb_filepath, "wb") as f: - f.write(output_graph_def.SerializeToString()) + import tf2onnx + with tf.device("/CPU:0"): + model_proto, _ = tf2onnx.convert._convert_common( + output_graph_def, + name='AMP', + input_names=['in_face:0','morph_value:0'], + output_names=['out_face_mask:0','out_celeb_face:0','out_celeb_face_mask:0'], + opset=13, + output_path=output_path) #override def get_model_filename_list(self): @@ -716,22 +699,24 @@ class AMPModel(ModelBase): src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) for i in range(bs): - self.last_src_samples_loss.append ( (target_src[i], target_srcm[i], target_srcm_em[i], src_loss[i] ) ) - self.last_dst_samples_loss.append ( (target_dst[i], target_dstm[i], target_dstm_em[i], dst_loss[i] ) ) + self.last_src_samples_loss.append ( (src_loss[i], warped_src[i], target_src[i], target_srcm[i], target_srcm_em[i]) ) + self.last_dst_samples_loss.append ( (dst_loss[i], warped_dst[i], target_dst[i], target_dstm[i], target_dstm_em[i]) ) if len(self.last_src_samples_loss) >= bs*16: - src_samples_loss = sorted(self.last_src_samples_loss, key=operator.itemgetter(3), reverse=True) - dst_samples_loss = sorted(self.last_dst_samples_loss, key=operator.itemgetter(3), reverse=True) + src_samples_loss = sorted(self.last_src_samples_loss, key=operator.itemgetter(0), reverse=True) + dst_samples_loss = sorted(self.last_dst_samples_loss, key=operator.itemgetter(0), reverse=True) - target_src = np.stack( [ x[0] for x in src_samples_loss[:bs] ] ) - target_srcm = np.stack( [ x[1] for x in src_samples_loss[:bs] ] ) - target_srcm_em = np.stack( [ x[2] for x in src_samples_loss[:bs] ] ) + warped_src = np.stack( [ x[1] for x in src_samples_loss[:bs] ] ) + target_src = np.stack( [ x[2] for x in src_samples_loss[:bs] ] ) + target_srcm = np.stack( [ x[3] for x in src_samples_loss[:bs] ] ) + target_srcm_em = np.stack( [ x[4] for x in src_samples_loss[:bs] ] ) - target_dst = np.stack( [ x[0] for x in dst_samples_loss[:bs] ] ) - target_dstm = np.stack( [ x[1] for x in dst_samples_loss[:bs] ] ) - target_dstm_em = np.stack( [ x[2] for x in dst_samples_loss[:bs] ] ) + warped_dst = np.stack( [ x[1] for x in dst_samples_loss[:bs] ] ) + target_dst = np.stack( [ x[2] for x in dst_samples_loss[:bs] ] ) + target_dstm = np.stack( [ x[3] for x in dst_samples_loss[:bs] ] ) + target_dstm_em = np.stack( [ x[4] for x in dst_samples_loss[:bs] ] ) - src_loss, dst_loss = self.src_dst_train (target_src, target_src, target_srcm, target_srcm_em, target_dst, target_dst, target_dstm, target_dstm_em) + src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) self.last_src_samples_loss = [] self.last_dst_samples_loss = [] @@ -741,7 +726,7 @@ class AMPModel(ModelBase): return ( ('src_loss', np.mean(src_loss) ), ('dst_loss', np.mean(dst_loss) ), ) #override - def onGetPreview(self, samples): + def onGetPreview(self, samples, for_history=False): ( (warped_src, target_src, target_srcm, target_srcm_em), (warped_dst, target_dst, target_dstm, target_dstm_em) ) = samples @@ -771,7 +756,7 @@ class AMPModel(ModelBase): result = [] - i = np.random.randint(n_samples) + i = np.random.randint(n_samples) if not for_history else 0 st = [ np.concatenate ((S[i], D[i], DD[i]*DDM_000[i]), axis=1) ] st += [ np.concatenate ((SS[i], DD[i], SD_075[i] ), axis=1) ] @@ -782,7 +767,6 @@ class AMPModel(ModelBase): st += [ np.concatenate ((SD_065[i], SD_075[i], SD_100[i]), axis=1) ] result += [ ('AMP morph list', np.concatenate (st, axis=0 )), ] - 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 )), ] @@ -791,7 +775,7 @@ class AMPModel(ModelBase): def predictor_func (self, face, morph_value): face = nn.to_data_format(face[None,...], self.model_data_format, "NHWC") - + bgr, mask_dst_dstm, mask_src_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face, morph_value) ] return bgr[0], mask_src_dstm[0][...,0], mask_dst_dstm[0][...,0] @@ -802,9 +786,9 @@ class AMPModel(ModelBase): def predictor_morph(face): return self.predictor_func(face, morph_factor) - - import merger + + import merger return predictor_morph, (self.options['resolution'], self.options['resolution'], 3), merger.MergerConfigMasked(face_type=self.face_type, default_mode = 'overlay') Model = AMPModel diff --git a/models/Model_Quick96/Model.py b/models/Model_Quick96/Model.py index fa139e5..fa9e215 100644 --- a/models/Model_Quick96/Model.py +++ b/models/Model_Quick96/Model.py @@ -278,7 +278,7 @@ class QModel(ModelBase): return ( ('src_loss', src_loss), ('dst_loss', dst_loss), ) #override - def onGetPreview(self, samples): + def onGetPreview(self, samples, for_history=False): ( (warped_src, target_src, target_srcm), (warped_dst, target_dst, target_dstm) ) = samples diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index b9b1c42..cc1f49a 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -659,11 +659,15 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.pretrain_just_disabled: self.update_sample_for_preview(force_new=True) - def dump_ckpt(self): + def dump_dflive (self): + output_path=self.get_strpath_storage_for_file('model.dflive') + + io.log_info(f'Dumping .dflive to {output_path}') + tf = nn.tf + nn.set_data_format('NCHW') - - with tf.device ('/CPU:0'): + with tf.device (nn.tf_default_device_name): warped_dst = tf.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face') warped_dst = tf.transpose(warped_dst, (0,3,1,2)) @@ -687,15 +691,26 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_pred_dst_dstm = tf.transpose(gpu_pred_dst_dstm, (0,2,3,1)) gpu_pred_src_dstm = tf.transpose(gpu_pred_src_dstm, (0,2,3,1)) - - saver = tf.train.Saver() tf.identity(gpu_pred_dst_dstm, name='out_face_mask') tf.identity(gpu_pred_src_dst, name='out_celeb_face') tf.identity(gpu_pred_src_dstm, name='out_celeb_face_mask') - saver.save(nn.tf_sess, self.get_strpath_storage_for_file('.ckpt') ) - + output_graph_def = tf.graph_util.convert_variables_to_constants( + nn.tf_sess, + tf.get_default_graph().as_graph_def(), + ['out_face_mask','out_celeb_face','out_celeb_face_mask'] + ) + import tf2onnx + with tf.device("/CPU:0"): + model_proto, _ = tf2onnx.convert._convert_common( + output_graph_def, + name='SAEHD', + input_names=['in_face:0'], + output_names=['out_face_mask:0','out_celeb_face:0','out_celeb_face_mask:0'], + opset=13, + output_path=output_path) + #override def get_model_filename_list(self): return self.model_filename_list @@ -751,7 +766,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... return ( ('src_loss', np.mean(src_loss) ), ('dst_loss', np.mean(dst_loss) ), ) #override - def onGetPreview(self, samples): + def onGetPreview(self, samples, for_history=False): ( (warped_src, target_src, target_srcm, target_srcm_em), (warped_dst, target_dst, target_dstm, target_dstm_em) ) = samples diff --git a/models/Model_XSeg/Model.py b/models/Model_XSeg/Model.py index 567a1f2..a27d5ad 100644 --- a/models/Model_XSeg/Model.py +++ b/models/Model_XSeg/Model.py @@ -164,7 +164,7 @@ class XSegModel(ModelBase): return ( ('loss', np.mean(loss) ), ) #override - def onGetPreview(self, samples): + def onGetPreview(self, samples, for_history=False): n_samples = min(4, self.get_batch_size(), 800 // self.resolution ) srcdst_samples, src_samples, dst_samples = samples diff --git a/requirements-colab.txt b/requirements-colab.txt index 814ed07..49e8a9f 100644 --- a/requirements-colab.txt +++ b/requirements-colab.txt @@ -6,4 +6,5 @@ ffmpeg-python==0.1.17 scikit-image==0.14.2 scipy==1.4.1 colorama -tensorflow-gpu==2.3.1 \ No newline at end of file +tensorflow-gpu==2.3.1 +tf2onnx==1.8.4 \ No newline at end of file diff --git a/requirements-cuda.txt b/requirements-cuda.txt index e31bbba..fbc5ba8 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -8,3 +8,4 @@ scipy==1.4.1 colorama tensorflow-gpu==2.4.0 pyqt5 +tf2onnx==1.8.4 \ No newline at end of file From 3c6bbe22b90c6a0f7bdd0b14ce38a3da5ddda350 Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 26 Jun 2021 10:55:33 +0400 Subject: [PATCH 03/30] AMP: added Inter dimensions param. The model is not changed. Should be equal or more than AutoEncoder dimensions. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU. --- models/Model_AMP/Model.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 5b75036..1588f56 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -36,6 +36,12 @@ class AMPModel(ModelBase): default_models_opt_on_gpu = self.options['models_opt_on_gpu'] = self.load_or_def_option('models_opt_on_gpu', True) default_ae_dims = self.options['ae_dims'] = self.load_or_def_option('ae_dims', 256) + + inter_dims = self.load_or_def_option('inter_dims', None) + if inter_dims is None: + inter_dims = self.options['ae_dims'] + default_inter_dims = self.options['inter_dims'] = inter_dims + default_e_dims = self.options['e_dims'] = self.load_or_def_option('e_dims', 64) default_d_dims = self.options['d_dims'] = self.options.get('d_dims', None) default_d_mask_dims = self.options['d_mask_dims'] = self.options.get('d_mask_dims', None) @@ -77,7 +83,9 @@ class AMPModel(ModelBase): default_d_mask_dims = self.options['d_mask_dims'] = self.load_or_def_option('d_mask_dims', default_d_mask_dims) if self.is_first_run(): - self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dimensions", default_ae_dims, add_info="32-1024", help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) + self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dimensions", default_ae_dims, add_info="32-1024", help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) + self.options['inter_dims'] = np.clip ( io.input_int("Inter dimensions", default_inter_dims, add_info="32-2048", help_message="Should be equal or more than AutoEncoder dimensions. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 2048 ) + e_dims = np.clip ( io.input_int("Encoder dimensions", default_e_dims, add_info="16-256", help_message="More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU." ), 16, 256 ) self.options['e_dims'] = e_dims + e_dims % 2 @@ -139,6 +147,7 @@ class AMPModel(ModelBase): input_ch=3 ae_dims = self.ae_dims = self.options['ae_dims'] + inter_dims = self.inter_dims = self.options['inter_dims'] e_dims = self.options['e_dims'] d_dims = self.options['d_dims'] d_mask_dims = self.options['d_mask_dims'] @@ -203,13 +212,13 @@ class AMPModel(ModelBase): def forward(self, inp): x = inp x = self.dense2(x) - x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, ae_dims) + x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, inter_dims) return x class Decoder(nn.ModelBase): def on_build(self ): - self.upscale0 = Upscale(ae_dims, d_dims*8, kernel_size=3) + self.upscale0 = Upscale(inter_dims, d_dims*8, kernel_size=3) self.upscale1 = Upscale(d_dims*8, d_dims*8, kernel_size=3) self.upscale2 = Upscale(d_dims*8, d_dims*4, kernel_size=3) self.upscale3 = Upscale(d_dims*4, d_dims*2, kernel_size=3) @@ -219,7 +228,7 @@ class AMPModel(ModelBase): self.res2 = ResidualBlock(d_dims*4, kernel_size=3) self.res3 = ResidualBlock(d_dims*2, kernel_size=3) - self.upscalem0 = Upscale(ae_dims, d_mask_dims*8, kernel_size=3) + self.upscalem0 = Upscale(inter_dims, d_mask_dims*8, kernel_size=3) self.upscalem1 = Upscale(d_mask_dims*8, d_mask_dims*8, kernel_size=3) self.upscalem2 = Upscale(d_mask_dims*8, d_mask_dims*4, kernel_size=3) self.upscalem3 = Upscale(d_mask_dims*4, d_mask_dims*2, kernel_size=3) @@ -402,9 +411,9 @@ class AMPModel(ModelBase): gpu_src_code = gpu_src_inter_src_code * inter_rnd_binomial + gpu_src_inter_dst_code * (1-inter_rnd_binomial) gpu_dst_code = gpu_dst_inter_dst_code - ae_dims_slice = tf.cast(ae_dims*self.morph_value_t[0], tf.int32) - gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, ae_dims_slice , lowest_dense_res, lowest_dense_res]), - tf.slice(gpu_dst_inter_dst_code, [0,ae_dims_slice,0,0], [-1,ae_dims-ae_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) + inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) + gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , lowest_dense_res, lowest_dense_res]), + tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,ae_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code) gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code) @@ -558,9 +567,9 @@ class AMPModel(ModelBase): gpu_dst_inter_src_code = self.inter_src ( gpu_dst_code) gpu_dst_inter_dst_code = self.inter_dst ( gpu_dst_code) - ae_dims_slice = tf.cast(ae_dims*self.morph_value_t[0], tf.int32) - gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, ae_dims_slice , lowest_dense_res, lowest_dense_res]), - tf.slice(gpu_dst_inter_dst_code, [0,ae_dims_slice,0,0], [-1,ae_dims-ae_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) + inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) + gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , lowest_dense_res, lowest_dense_res]), + tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,ae_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code) _, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code) From 714234bce40e60b1d68b4a3d1ffa25cf3436b0a3 Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 26 Jun 2021 18:03:04 +0400 Subject: [PATCH 04/30] missing file --- mainscripts/DumpDFLive.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 mainscripts/DumpDFLive.py diff --git a/mainscripts/DumpDFLive.py b/mainscripts/DumpDFLive.py new file mode 100644 index 0000000..704cf23 --- /dev/null +++ b/mainscripts/DumpDFLive.py @@ -0,0 +1,22 @@ +import os +import sys +import traceback +import queue +import threading +import time +import numpy as np +import itertools +from pathlib import Path +from core import pathex +from core import imagelib +import cv2 +import models +from core.interact import interact as io + + +def main(model_class_name, saved_models_path): + model = models.import_model(model_class_name)( + is_training=False, + saved_models_path=saved_models_path, + cpu_only=True) + model.dump_dflive () From 90a74efd89a64861aa684759edf105a39dba2f39 Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 26 Jun 2021 23:06:30 +0400 Subject: [PATCH 05/30] fix --- models/Model_AMP/Model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 1588f56..2c60741 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -207,7 +207,7 @@ class AMPModel(ModelBase): class Inter(nn.ModelBase): def on_build(self): - self.dense2 = nn.Dense( ae_dims, lowest_dense_res * lowest_dense_res * ae_dims ) + self.dense2 = nn.Dense( ae_dims, lowest_dense_res * lowest_dense_res * inter_dims ) def forward(self, inp): x = inp From c43b3b161b606f418a3e4d283a384588c53b748f Mon Sep 17 00:00:00 2001 From: iperov Date: Sun, 27 Jun 2021 20:43:44 +0400 Subject: [PATCH 06/30] _ --- mainscripts/FacesetResizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mainscripts/FacesetResizer.py b/mainscripts/FacesetResizer.py index acc4ea4..4bcd1b8 100644 --- a/mainscripts/FacesetResizer.py +++ b/mainscripts/FacesetResizer.py @@ -166,7 +166,7 @@ class FacesetResizerSubprocessor(Subprocessor): def process_folder ( dirpath): - image_size = io.input_int(f"New image size", 512, valid_range=[256,2048]) + image_size = io.input_int(f"New image size", 512, valid_range=[128,2048]) face_type = io.input_str ("Change face type", 'same', ['h','mf','f','wf','head','same']).lower() if face_type == 'same': From 9d6b6feb1fb278b80e6490aa559ff5e636fc7c67 Mon Sep 17 00:00:00 2001 From: iperov Date: Mon, 28 Jun 2021 14:27:09 +0400 Subject: [PATCH 07/30] _ --- models/Model_AMP/Model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 2c60741..fa093fe 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -413,7 +413,7 @@ class AMPModel(ModelBase): inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , lowest_dense_res, lowest_dense_res]), - tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,ae_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) + tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code) gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code) @@ -569,7 +569,7 @@ class AMPModel(ModelBase): inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , lowest_dense_res, lowest_dense_res]), - tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,ae_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) + tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code) _, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code) From 63c794b3d091e7df5eeaea62a5f414babc5bd3f4 Mon Sep 17 00:00:00 2001 From: iperov Date: Mon, 28 Jun 2021 18:26:54 +0400 Subject: [PATCH 08/30] rename dump_dflive to export_dfm --- main.py | 10 +++++----- mainscripts/{DumpDFLive.py => ExportDFM.py} | 2 +- models/Model_AMP/Model.py | 6 +++--- models/Model_SAEHD/Model.py | 6 +++--- 4 files changed, 12 insertions(+), 12 deletions(-) rename mainscripts/{DumpDFLive.py => ExportDFM.py} (95%) diff --git a/main.py b/main.py index 04337af..493c7ed 100644 --- a/main.py +++ b/main.py @@ -148,15 +148,15 @@ if __name__ == "__main__": p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+') p.set_defaults (func=process_train) - def process_dumpdflive(arguments): + def process_exportdfm(arguments): osex.set_process_lowest_prio() - from mainscripts import DumpDFLive - DumpDFLive.main(model_class_name = arguments.model_name, saved_models_path = Path(arguments.model_dir)) + from mainscripts import ExportDFM + ExportDFM.main(model_class_name = arguments.model_name, saved_models_path = Path(arguments.model_dir)) - p = subparsers.add_parser( "dumpdflive", help="Dump model to use in DFLive.") + p = subparsers.add_parser( "exportdfm", help="Export model to use in DeepFaceLive.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.") p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.") - p.set_defaults (func=process_dumpdflive) + p.set_defaults (func=process_exportdfm) def process_merge(arguments): osex.set_process_lowest_prio() diff --git a/mainscripts/DumpDFLive.py b/mainscripts/ExportDFM.py similarity index 95% rename from mainscripts/DumpDFLive.py rename to mainscripts/ExportDFM.py index 704cf23..2b0f978 100644 --- a/mainscripts/DumpDFLive.py +++ b/mainscripts/ExportDFM.py @@ -19,4 +19,4 @@ def main(model_class_name, saved_models_path): is_training=False, saved_models_path=saved_models_path, cpu_only=True) - model.dump_dflive () + model.export_dfm () diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index fa093fe..b6732dc 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -638,10 +638,10 @@ class AMPModel(ModelBase): self.update_sample_for_preview(force_new=True) - def dump_dflive (self): - output_path=self.get_strpath_storage_for_file('model.dflive') + def export_dfm (self): + output_path=self.get_strpath_storage_for_file('model.dfm') - io.log_info(f'Dumping .dflive to {output_path}') + io.log_info(f'Dumping .dfm to {output_path}') tf = nn.tf with tf.device (nn.tf_default_device_name): diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index cc1f49a..7e548db 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -659,10 +659,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.pretrain_just_disabled: self.update_sample_for_preview(force_new=True) - def dump_dflive (self): - output_path=self.get_strpath_storage_for_file('model.dflive') + def export_dfm (self): + output_path=self.get_strpath_storage_for_file('model.dfm') - io.log_info(f'Dumping .dflive to {output_path}') + io.log_info(f'Dumping .dfm to {output_path}') tf = nn.tf nn.set_data_format('NCHW') From 075422e0c59879c0dc366b4bffce68c141b55d00 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 1 Jul 2021 14:25:54 +0400 Subject: [PATCH 09/30] amp: fixes, removed pretrain option --- models/Model_AMP/Model.py | 140 ++++++++++++++++++++------------------ 1 file changed, 72 insertions(+), 68 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index b6732dc..1f7f505 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -57,7 +57,7 @@ class AMPModel(ModelBase): default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True) default_ct_mode = self.options['ct_mode'] = self.load_or_def_option('ct_mode', 'none') 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_pretrain = self.options['pretrain'] = self.load_or_def_option('pretrain', False) ask_override = self.ask_override() @@ -73,7 +73,7 @@ class AMPModel(ModelBase): resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 32 .") resolution = np.clip ( (resolution // 32) * 32, min_res, max_res) self.options['resolution'] = resolution - self.options['face_type'] = io.input_str ("Face type", default_face_type, ['wf','head'], help_message="whole face / head").lower() + self.options['face_type'] = io.input_str ("Face type", default_face_type, ['f','wf','head'], help_message="whole face / head").lower() default_d_dims = self.options['d_dims'] = self.load_or_def_option('d_dims', 64) @@ -86,7 +86,6 @@ class AMPModel(ModelBase): self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dimensions", default_ae_dims, add_info="32-1024", help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) self.options['inter_dims'] = np.clip ( io.input_int("Inter dimensions", default_inter_dims, add_info="32-2048", help_message="Should be equal or more than AutoEncoder dimensions. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 2048 ) - e_dims = np.clip ( io.input_int("Encoder dimensions", default_e_dims, add_info="16-256", help_message="More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU." ), 16, 256 ) self.options['e_dims'] = e_dims + e_dims % 2 @@ -130,10 +129,10 @@ class AMPModel(ModelBase): self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") self.options['clipgrad'] = io.input_bool ("Enable gradient clipping", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.") - self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly. Forces random_warp=N, random_flips=Y, gan_power=0.0, lr_dropout=N, uniform_yaw=Y") + #self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly. Forces random_warp=N, random_flips=Y, gan_power=0.0, lr_dropout=N, uniform_yaw=Y") self.gan_model_changed = (default_gan_patch_size != self.options['gan_patch_size']) or (default_gan_dims != self.options['gan_dims']) - self.pretrain_just_disabled = (default_pretrain == True and self.options['pretrain'] == False) + #self.pretrain_just_disabled = (default_pretrain == True and self.options['pretrain'] == False) #override def on_initialize(self): @@ -151,7 +150,7 @@ class AMPModel(ModelBase): e_dims = self.options['e_dims'] d_dims = self.options['d_dims'] d_mask_dims = self.options['d_mask_dims'] - lowest_dense_res = self.lowest_dense_res = resolution // 32 + inter_res = self.inter_res = resolution // 32 class Downscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=5 ): @@ -189,7 +188,7 @@ class AMPModel(ModelBase): self.down4 = Downscale(e_dims*4, e_dims*8, kernel_size=5) self.down5 = Downscale(e_dims*8, e_dims*8, kernel_size=5) self.res5 = ResidualBlock(e_dims*8) - self.dense1 = nn.Dense( lowest_dense_res*lowest_dense_res*e_dims*8, ae_dims ) + self.dense1 = nn.Dense( (( resolution//(2**5) )**2) * e_dims*8, ae_dims ) def forward(self, inp): x = inp @@ -207,12 +206,12 @@ class AMPModel(ModelBase): class Inter(nn.ModelBase): def on_build(self): - self.dense2 = nn.Dense( ae_dims, lowest_dense_res * lowest_dense_res * inter_dims ) + self.dense2 = nn.Dense(ae_dims, inter_res * inter_res * inter_dims) def forward(self, inp): x = inp x = self.dense2(x) - x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, inter_dims) + x = nn.reshape_4D (x, inter_res, inter_res, inter_dims) return x @@ -265,7 +264,8 @@ class AMPModel(ModelBase): m = tf.nn.sigmoid(self.out_convm(m)) return x, m - self.face_type = {'wf' : FaceType.WHOLE_FACE, + self.face_type = {'f' : FaceType.FULL, + 'wf' : FaceType.WHOLE_FACE, 'head' : FaceType.HEAD}[ self.options['face_type'] ] if 'eyes_prio' in self.options: @@ -276,20 +276,24 @@ class AMPModel(ModelBase): morph_factor = self.options['morph_factor'] - pretrain = self.pretrain = self.options['pretrain'] - if self.pretrain_just_disabled: - self.set_iter(0) + gan_power = self.gan_power = self.options['gan_power'] + random_warp = self.options['random_warp'] + random_src_flip = self.random_src_flip + random_dst_flip = self.random_dst_flip + + #pretrain = self.pretrain = self.options['pretrain'] + #if self.pretrain_just_disabled: + # self.set_iter(0) + # self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power'] + # random_warp = False if self.pretrain else self.options['random_warp'] + # random_src_flip = self.random_src_flip if not self.pretrain else True + # random_dst_flip = self.random_dst_flip if not self.pretrain else True - self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power'] - random_warp = False if self.pretrain else self.options['random_warp'] - random_src_flip = self.random_src_flip if not self.pretrain else True - random_dst_flip = self.random_dst_flip if not self.pretrain else True - - if self.pretrain: - self.options_show_override['gan_power'] = 0.0 - self.options_show_override['random_warp'] = False - self.options_show_override['lr_dropout'] = 'n' - self.options_show_override['uniform_yaw'] = True + # if self.pretrain: + # self.options_show_override['gan_power'] = 0.0 + # self.options_show_override['random_warp'] = False + # self.options_show_override['lr_dropout'] = 'n' + # self.options_show_override['uniform_yaw'] = True masked_training = self.options['masked_training'] ct_mode = self.options['ct_mode'] @@ -340,15 +344,15 @@ class AMPModel(ModelBase): # Initialize optimizers lr=5e-5 - lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] and not self.pretrain else 1.0 + lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] else 1.0 #and not self.pretrain clipnorm = 1.0 if self.options['clipgrad'] else 0.0 - self.all_weights = self.encoder.get_weights() + self.inter_src.get_weights() + self.inter_dst.get_weights() + self.decoder.get_weights() - if pretrain: - self.trainable_weights = self.encoder.get_weights() + self.inter_dst.get_weights() + self.decoder.get_weights() - else: - self.trainable_weights = self.encoder.get_weights() + self.inter_src.get_weights() + self.inter_dst.get_weights() + self.decoder.get_weights() + self.all_weights = self.encoder.get_weights() + self.decoder.get_weights() + #if pretrain: + # self.trainable_weights = self.encoder.get_weights() + self.decoder.get_weights() + #else: + self.trainable_weights = self.encoder.get_weights() + self.decoder.get_weights() self.src_dst_opt = nn.AdaBelief(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt') self.src_dst_opt.initialize_variables (self.all_weights, vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu') @@ -396,24 +400,24 @@ class AMPModel(ModelBase): gpu_src_code = self.encoder (gpu_warped_src) gpu_dst_code = self.encoder (gpu_warped_dst) - if pretrain: - gpu_src_inter_src_code = self.inter_src (gpu_src_code) - gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) - gpu_src_code = gpu_src_inter_src_code * nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor) - gpu_dst_code = gpu_src_dst_code = gpu_dst_inter_dst_code * nn.random_binomial( [bs_per_gpu, gpu_dst_inter_dst_code.shape.as_list()[1], 1,1] , p=0.25) - else: - gpu_src_inter_src_code = self.inter_src (gpu_src_code) - gpu_src_inter_dst_code = self.inter_dst (gpu_src_code) - gpu_dst_inter_src_code = self.inter_src (gpu_dst_code) - gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) + # if pretrain: + # gpu_src_inter_src_code = self.inter_src (gpu_src_code) + # gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) + # gpu_src_code = gpu_src_inter_src_code * nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor) + # gpu_dst_code = gpu_src_dst_code = gpu_dst_inter_dst_code * nn.random_binomial( [bs_per_gpu, gpu_dst_inter_dst_code.shape.as_list()[1], 1,1] , p=morph_factor) + # else: + gpu_src_inter_src_code = self.inter_src (gpu_src_code) + gpu_src_inter_dst_code = self.inter_dst (gpu_src_code) + gpu_dst_inter_src_code = self.inter_src (gpu_dst_code) + gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) - inter_rnd_binomial = nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor) - gpu_src_code = gpu_src_inter_src_code * inter_rnd_binomial + gpu_src_inter_dst_code * (1-inter_rnd_binomial) - gpu_dst_code = gpu_dst_inter_dst_code + inter_rnd_binomial = nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor) + gpu_src_code = gpu_src_inter_src_code * inter_rnd_binomial + gpu_src_inter_dst_code * (1-inter_rnd_binomial) + gpu_dst_code = gpu_dst_inter_dst_code - inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) - gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , lowest_dense_res, lowest_dense_res]), - tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) + inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) + gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , inter_res, inter_res]), + tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, inter_res,inter_res]) ), 1 ) gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code) gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code) @@ -467,15 +471,15 @@ class AMPModel(ModelBase): gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_em - gpu_pred_src_src*gpu_target_srcm_em ), axis=[1,2,3]) gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) - else: - gpu_src_loss = gpu_dst_loss + #else: + # gpu_src_loss = gpu_dst_loss gpu_src_losses += [gpu_src_loss] - if pretrain: - gpu_G_loss = gpu_dst_loss - else: - gpu_G_loss = gpu_src_loss + gpu_dst_loss + #if pretrain: + # gpu_G_loss = gpu_dst_loss + #else: + gpu_G_loss = gpu_src_loss + gpu_dst_loss def DLossOnes(logits): return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(logits), logits=logits), axis=[1,2,3]) @@ -568,8 +572,8 @@ class AMPModel(ModelBase): gpu_dst_inter_dst_code = self.inter_dst ( gpu_dst_code) inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) - gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , lowest_dense_res, lowest_dense_res]), - tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 ) + gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , inter_res, inter_res]), + tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, inter_res,inter_res]) ), 1 ) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code) _, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code) @@ -581,15 +585,15 @@ class AMPModel(ModelBase): # Loading/initializing all models/optimizers weights for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"): - if self.pretrain_just_disabled: - do_init = False - if model == self.inter_src or model == self.inter_dst: + # if self.pretrain_just_disabled: + # do_init = False + # if model == self.inter_src or model == self.inter_dst: + # do_init = True + # else: + do_init = self.is_first_run() + if self.is_training and gan_power != 0 and model == self.GAN: + if self.gan_model_changed: do_init = True - else: - do_init = self.is_first_run() - if self.is_training and gan_power != 0 and model == self.GAN: - if self.gan_model_changed: - do_init = True if not do_init: do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) ) @@ -599,10 +603,10 @@ class AMPModel(ModelBase): # initializing sample generators if self.is_training: - training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path() - training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path() + training_data_src_path = self.training_data_src_path #if not self.pretrain else self.get_pretraining_data_path() + training_data_dst_path = self.training_data_dst_path #if not self.pretrain else self.get_pretraining_data_path() - random_ct_samples_path=training_data_dst_path if ct_mode is not None and not self.pretrain else None + random_ct_samples_path=training_data_dst_path if ct_mode is not None else None #and not self.pretrain cpu_count = min(multiprocessing.cpu_count(), 8) src_generators_count = cpu_count // 2 @@ -618,7 +622,7 @@ class AMPModel(ModelBase): {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, ], - uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain, + uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain, generators_count=src_generators_count ), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(), @@ -628,14 +632,14 @@ class AMPModel(ModelBase): {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, ], - uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain, + uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain, generators_count=dst_generators_count ) ]) self.last_src_samples_loss = [] self.last_dst_samples_loss = [] - if self.pretrain_just_disabled: - self.update_sample_for_preview(force_new=True) + #if self.pretrain_just_disabled: + # self.update_sample_for_preview(force_new=True) def export_dfm (self): From 0444a9a5653c2d9dcfbe311eb3f79b16f63c674e Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 1 Jul 2021 14:29:11 +0400 Subject: [PATCH 10/30] default morph factor is now 0.1 --- models/Model_AMP/Model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 1f7f505..ca3e1bd 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -45,7 +45,7 @@ class AMPModel(ModelBase): default_e_dims = self.options['e_dims'] = self.load_or_def_option('e_dims', 64) default_d_dims = self.options['d_dims'] = self.options.get('d_dims', None) default_d_mask_dims = self.options['d_mask_dims'] = self.options.get('d_mask_dims', None) - default_morph_factor = self.options['morph_factor'] = self.options.get('morph_factor', 0.33) + default_morph_factor = self.options['morph_factor'] = self.options.get('morph_factor', 0.1) default_masked_training = self.options['masked_training'] = self.load_or_def_option('masked_training', True) default_eyes_mouth_prio = self.options['eyes_mouth_prio'] = self.load_or_def_option('eyes_mouth_prio', True) default_uniform_yaw = self.options['uniform_yaw'] = self.load_or_def_option('uniform_yaw', False) From 446b95942fd6c4ced05a825ff758a3492d34d219 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 1 Jul 2021 15:08:11 +0400 Subject: [PATCH 11/30] fix --- models/Model_AMP/Model.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index ca3e1bd..62468ae 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -459,18 +459,18 @@ class AMPModel(ModelBase): gpu_dst_loss += 0.1*tf.reduce_mean(tf.square(gpu_pred_dst_dst_anti_masked-gpu_target_dst_anti_masked),axis=[1,2,3] ) gpu_dst_losses += [gpu_dst_loss] - if not pretrain: - if resolution < 256: - gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) - else: - gpu_src_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) - gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) - gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3]) + #if not pretrain: + if resolution < 256: + gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + else: + gpu_src_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) + gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3]) - if eyes_mouth_prio: - gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_em - gpu_pred_src_src*gpu_target_srcm_em ), axis=[1,2,3]) + if eyes_mouth_prio: + gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_em - gpu_pred_src_src*gpu_target_srcm_em ), axis=[1,2,3]) - gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) + gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) #else: # gpu_src_loss = gpu_dst_loss From 9e0bd342532713d02973800d5d6c8203971e9b03 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 1 Jul 2021 15:50:51 +0400 Subject: [PATCH 12/30] update torrent link --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 529ebbc..83bbb46 100644 --- a/README.md +++ b/README.md @@ -194,7 +194,7 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul -Windows (magnet link) +Windows (magnet link) Last release. Use torrent client to download. From b4b72d056f26b7b51747d9559eb252fe3c8b57f7 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 1 Jul 2021 21:58:23 +0400 Subject: [PATCH 13/30] upd torrent link --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 83bbb46..d6a219f 100644 --- a/README.md +++ b/README.md @@ -194,7 +194,7 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul -Windows (magnet link) +Windows (magnet link) Last release. Use torrent client to download. From ee1bc83a14aee303751ef87e2a87b02c09a9f9dd Mon Sep 17 00:00:00 2001 From: iperov Date: Wed, 7 Jul 2021 21:45:43 +0400 Subject: [PATCH 14/30] AMP, SAEHD: reverted GAN to december version. --- core/leras/models/PatchDiscriminator.py | 18 ++---------------- models/Model_AMP/Model.py | 2 +- models/Model_SAEHD/Model.py | 4 ++-- 3 files changed, 5 insertions(+), 19 deletions(-) diff --git a/core/leras/models/PatchDiscriminator.py b/core/leras/models/PatchDiscriminator.py index 343e000..868151d 100644 --- a/core/leras/models/PatchDiscriminator.py +++ b/core/leras/models/PatchDiscriminator.py @@ -146,11 +146,7 @@ class UNetPatchDiscriminator(nn.ModelBase): prev_ch = in_ch self.convs = [] - self.res1 = [] - self.res2 = [] self.upconvs = [] - self.upres1 = [] - self.upres2 = [] layers = self.find_archi(patch_size) level_chs = { i-1:v for i,v in enumerate([ min( base_ch * (2**i), 512 ) for i in range(len(layers)+1)]) } @@ -160,14 +156,8 @@ class UNetPatchDiscriminator(nn.ModelBase): for i, (kernel_size, strides) in enumerate(layers): self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME') ) - self.res1.append ( ResidualBlock(level_chs[i]) ) - self.res2.append ( ResidualBlock(level_chs[i]) ) - self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME') ) - self.upres1.insert (0, ResidualBlock(level_chs[i-1]*2) ) - self.upres2.insert (0, ResidualBlock(level_chs[i-1]*2) ) - self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID') self.center_out = nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID') @@ -178,19 +168,15 @@ class UNetPatchDiscriminator(nn.ModelBase): x = tf.nn.leaky_relu( self.in_conv(x), 0.2 ) encs = [] - for conv, res1,res2 in zip(self.convs, self.res1, self.res2): + for conv in self.convs: encs.insert(0, x) x = tf.nn.leaky_relu( conv(x), 0.2 ) - x = res1(x) - x = res2(x) center_out, x = self.center_out(x), tf.nn.leaky_relu( self.center_conv(x), 0.2 ) - for i, (upconv, enc, upres1, upres2 ) in enumerate(zip(self.upconvs, encs, self.upres1, self.upres2)): + for i, (upconv, enc) in enumerate(zip(self.upconvs, encs)): x = tf.nn.leaky_relu( upconv(x), 0.2 ) x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis) - x = upres1(x) - x = upres2(x) return center_out, self.out_conv(x) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 62468ae..25149f4 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -123,7 +123,7 @@ class AMPModel(ModelBase): gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 ) self.options['gan_patch_size'] = gan_patch_size - gan_dims = np.clip ( io.input_int("GAN dimensions", default_gan_dims, add_info="4-64", help_message="The dimensions of the GAN network. The higher dimensions, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is 16." ), 4, 64 ) + gan_dims = np.clip ( io.input_int("GAN dimensions", default_gan_dims, add_info="4-512", help_message="The dimensions of the GAN network. The higher dimensions, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is 16." ), 4, 512 ) self.options['gan_dims'] = gan_dims self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 7e548db..147defe 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -156,7 +156,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 ) self.options['gan_patch_size'] = gan_patch_size - gan_dims = np.clip ( io.input_int("GAN dimensions", default_gan_dims, add_info="4-64", help_message="The dimensions of the GAN network. The higher dimensions, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is 16." ), 4, 64 ) + gan_dims = np.clip ( io.input_int("GAN dimensions", default_gan_dims, add_info="4-512", help_message="The dimensions of the GAN network. The higher dimensions, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is 16." ), 4, 512 ) self.options['gan_dims'] = gan_dims if 'df' in self.options['archi']: @@ -467,7 +467,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... gpu_G_loss += self.options['true_face_power']*DLoss(gpu_src_code_d_ones, gpu_src_code_d) - gpu_D_code_loss = (DLoss(gpu_src_code_d_ones , gpu_dst_code_d) + \ + gpu_D_code_loss = (DLoss(gpu_dst_code_d_ones , gpu_dst_code_d) + \ DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5 gpu_D_code_loss_gvs += [ nn.gradients (gpu_D_code_loss, self.code_discriminator.get_weights() ) ] From 35877dbfd724c22040f421e93c1adbb7142e5b5d Mon Sep 17 00:00:00 2001 From: iperov Date: Wed, 14 Jul 2021 23:03:54 +0400 Subject: [PATCH 15/30] AMP, SAEHD: added option use_fp16: Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters. AMP: removed lr_dropout, now it is enabled by default; --- core/leras/archis/DeepFakeArchi.py | 97 +++++++++++-------------- core/leras/models/PatchDiscriminator.py | 33 ++++++--- models/Model_AMP/Model.py | 60 ++++++++------- models/Model_SAEHD/Model.py | 10 ++- 4 files changed, 102 insertions(+), 98 deletions(-) diff --git a/core/leras/archis/DeepFakeArchi.py b/core/leras/archis/DeepFakeArchi.py index 9818ab8..38b3f87 100644 --- a/core/leras/archis/DeepFakeArchi.py +++ b/core/leras/archis/DeepFakeArchi.py @@ -8,12 +8,15 @@ class DeepFakeArchi(nn.ArchiBase): mod None - default 'quick' """ - def __init__(self, resolution, mod=None, opts=None): + def __init__(self, resolution, use_fp16=False, mod=None, opts=None): super().__init__() if opts is None: opts = '' + + conv_dtype = tf.float16 if use_fp16 else tf.float32 + if mod is None: class Downscale(nn.ModelBase): def __init__(self, in_ch, out_ch, kernel_size=5, *kwargs ): @@ -23,7 +26,7 @@ class DeepFakeArchi(nn.ArchiBase): super().__init__(*kwargs) def on_build(self, *args, **kwargs ): - self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME') + self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME', dtype=conv_dtype) def forward(self, x): x = self.conv1(x) @@ -40,7 +43,7 @@ class DeepFakeArchi(nn.ArchiBase): last_ch = in_ch for i in range(n_downscales): cur_ch = ch*( min(2**i, 8) ) - self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size) ) + self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size)) last_ch = self.downs[-1].get_out_ch() def forward(self, inp): @@ -50,8 +53,8 @@ class DeepFakeArchi(nn.ArchiBase): return x class Upscale(nn.ModelBase): - def on_build(self, in_ch, out_ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME') + def on_build(self, in_ch, out_ch, kernel_size=3): + self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) def forward(self, x): x = self.conv1(x) @@ -60,9 +63,9 @@ class DeepFakeArchi(nn.ArchiBase): return x class ResidualBlock(nn.ModelBase): - def on_build(self, ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME') - self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME') + def on_build(self, ch, kernel_size=3): + self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) + self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) def forward(self, inp): x = self.conv1(inp) @@ -80,8 +83,13 @@ class DeepFakeArchi(nn.ArchiBase): def on_build(self): self.down1 = DownscaleBlock(self.in_ch, self.e_ch, n_downscales=4, kernel_size=5) - def forward(self, inp): - return nn.flatten(self.down1(inp)) + def forward(self, x): + if use_fp16: + x = tf.cast(x, tf.float16) + x = nn.flatten(self.down1(x)) + if use_fp16: + x = tf.cast(x, tf.float32) + return x def get_out_res(self, res): return res // (2**4) @@ -98,9 +106,10 @@ class DeepFakeArchi(nn.ArchiBase): def on_build(self): in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch + if 'u' in opts: self.dense_norm = nn.DenseNorm() - + self.dense1 = nn.Dense( in_ch, ae_ch ) self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch ) self.upscale1 = Upscale(ae_out_ch, ae_out_ch) @@ -112,6 +121,9 @@ class DeepFakeArchi(nn.ArchiBase): x = self.dense1(x) x = self.dense2(x) x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch) + + if use_fp16: + x = tf.cast(x, tf.float16) x = self.upscale1(x) return x @@ -122,7 +134,7 @@ class DeepFakeArchi(nn.ArchiBase): return self.ae_out_ch class Decoder(nn.ModelBase): - def on_build(self, in_ch, d_ch, d_mask_ch ): + def on_build(self, in_ch, d_ch, d_mask_ch): self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3) self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3) self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3) @@ -131,25 +143,23 @@ class DeepFakeArchi(nn.ArchiBase): self.res1 = ResidualBlock(d_ch*4, kernel_size=3) self.res2 = ResidualBlock(d_ch*2, kernel_size=3) - self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME') + self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype) self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3) self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3) self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3) - self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME') + self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', dtype=conv_dtype) if 'd' in opts: - self.out_conv1 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME') - self.out_conv2 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME') - self.out_conv3 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME') + self.out_conv1 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) + self.out_conv2 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) + self.out_conv3 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) self.upscalem3 = Upscale(d_mask_ch*2, d_mask_ch*1, kernel_size=3) - self.out_convm = nn.Conv2D( d_mask_ch*1, 1, kernel_size=1, padding='SAME') + self.out_convm = nn.Conv2D( d_mask_ch*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype) else: - self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME') - - def forward(self, inp): - z = inp + self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', dtype=conv_dtype) + def forward(self, z): x = self.upscale0(z) x = self.res0(x) x = self.upscale1(x) @@ -157,40 +167,11 @@ class DeepFakeArchi(nn.ArchiBase): x = self.upscale2(x) x = self.res2(x) - if 'd' in opts: - x0 = tf.nn.sigmoid(self.out_conv(x)) - x0 = nn.upsample2d(x0) - x1 = tf.nn.sigmoid(self.out_conv1(x)) - x1 = nn.upsample2d(x1) - x2 = tf.nn.sigmoid(self.out_conv2(x)) - x2 = nn.upsample2d(x2) - x3 = tf.nn.sigmoid(self.out_conv3(x)) - x3 = nn.upsample2d(x3) - - if nn.data_format == "NHWC": - tile_cfg = ( 1, resolution // 2, resolution //2, 1) - else: - tile_cfg = ( 1, 1, resolution // 2, resolution //2 ) - - z0 = tf.concat ( ( tf.concat ( ( tf.ones ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ), - tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] ) - - z0 = tf.tile ( z0, tile_cfg ) - - z1 = tf.concat ( ( tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.ones ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ), - tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] ) - z1 = tf.tile ( z1, tile_cfg ) - - z2 = tf.concat ( ( tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ), - tf.concat ( ( tf.ones ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] ) - z2 = tf.tile ( z2, tile_cfg ) - - z3 = tf.concat ( ( tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.zeros ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ), - tf.concat ( ( tf.zeros ( (1,1,1,1) ), tf.ones ( (1,1,1,1) ) ), axis=nn.conv2d_spatial_axes[1] ) ), axis=nn.conv2d_spatial_axes[0] ) - z3 = tf.tile ( z3, tile_cfg ) - - x = x0*z0 + x1*z1 + x2*z2 + x3*z3 + x = tf.nn.sigmoid( nn.depth_to_space(tf.concat( (self.out_conv(x), + self.out_conv1(x), + self.out_conv2(x), + self.out_conv3(x)), nn.conv2d_ch_axis), 2) ) else: x = tf.nn.sigmoid(self.out_conv(x)) @@ -201,7 +182,11 @@ class DeepFakeArchi(nn.ArchiBase): if 'd' in opts: m = self.upscalem3(m) m = tf.nn.sigmoid(self.out_convm(m)) - + + if use_fp16: + x = tf.cast(x, tf.float32) + m = tf.cast(m, tf.float32) + return x, m self.Encoder = Encoder diff --git a/core/leras/models/PatchDiscriminator.py b/core/leras/models/PatchDiscriminator.py index 868151d..9b94e9f 100644 --- a/core/leras/models/PatchDiscriminator.py +++ b/core/leras/models/PatchDiscriminator.py @@ -130,12 +130,14 @@ class UNetPatchDiscriminator(nn.ModelBase): q=x[np.abs(np.array(x)-target_patch_size).argmin()] return s[q][2] - def on_build(self, patch_size, in_ch, base_ch = 16): - + def on_build(self, patch_size, in_ch, base_ch = 16, use_fp16 = False): + self.use_fp16 = use_fp16 + conv_dtype = tf.float16 if use_fp16 else tf.float32 + class ResidualBlock(nn.ModelBase): def on_build(self, ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME') - self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME') + self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) + self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) def forward(self, inp): x = self.conv1(inp) @@ -151,20 +153,23 @@ class UNetPatchDiscriminator(nn.ModelBase): level_chs = { i-1:v for i,v in enumerate([ min( base_ch * (2**i), 512 ) for i in range(len(layers)+1)]) } - self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID') + self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID', dtype=conv_dtype) for i, (kernel_size, strides) in enumerate(layers): - self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME') ) + self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) ) - self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME') ) + self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) ) - self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID') + self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID', dtype=conv_dtype) - self.center_out = nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID') - self.center_conv = nn.Conv2D( level_chs[len(layers)-1], level_chs[len(layers)-1], kernel_size=1, padding='VALID') + self.center_out = nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID', dtype=conv_dtype) + self.center_conv = nn.Conv2D( level_chs[len(layers)-1], level_chs[len(layers)-1], kernel_size=1, padding='VALID', dtype=conv_dtype) def forward(self, x): + if self.use_fp16: + x = tf.cast(x, tf.float16) + x = tf.nn.leaky_relu( self.in_conv(x), 0.2 ) encs = [] @@ -178,6 +183,12 @@ class UNetPatchDiscriminator(nn.ModelBase): x = tf.nn.leaky_relu( upconv(x), 0.2 ) x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis) - return center_out, self.out_conv(x) + x = self.out_conv(x) + + if self.use_fp16: + center_out = tf.cast(center_out, tf.float32) + x = tf.cast(x, tf.float32) + + return center_out, x nn.UNetPatchDiscriminator = UNetPatchDiscriminator diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 25149f4..f6129ad 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -30,7 +30,8 @@ class AMPModel(ModelBase): yn_str = {True:'y',False:'n'} min_res = 64 max_res = 640 - + + default_usefp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) default_resolution = self.options['resolution'] = self.load_or_def_option('resolution', 224) default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'wf') default_models_opt_on_gpu = self.options['models_opt_on_gpu'] = self.load_or_def_option('models_opt_on_gpu', True) @@ -50,10 +51,6 @@ class AMPModel(ModelBase): default_eyes_mouth_prio = self.options['eyes_mouth_prio'] = self.load_or_def_option('eyes_mouth_prio', True) default_uniform_yaw = self.options['uniform_yaw'] = self.load_or_def_option('uniform_yaw', False) - lr_dropout = self.load_or_def_option('lr_dropout', 'n') - lr_dropout = {True:'y', False:'n'}.get(lr_dropout, lr_dropout) #backward comp - default_lr_dropout = self.options['lr_dropout'] = lr_dropout - default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True) default_ct_mode = self.options['ct_mode'] = self.load_or_def_option('ct_mode', 'none') default_clipgrad = self.options['clipgrad'] = self.load_or_def_option('clipgrad', False) @@ -68,7 +65,8 @@ class AMPModel(ModelBase): self.ask_random_src_flip() self.ask_random_dst_flip() self.ask_batch_size(suggest_batch_size) - + self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters.') + if self.is_first_run(): resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 32 .") resolution = np.clip ( (resolution // 32) * 32, min_res, max_res) @@ -113,11 +111,9 @@ class AMPModel(ModelBase): if self.is_first_run() or ask_override: self.options['models_opt_on_gpu'] = io.input_bool ("Place models and optimizer on GPU", default_models_opt_on_gpu, help_message="When you train on one GPU, by default model and optimizer weights are placed on GPU to accelerate the process. You can place they on CPU to free up extra VRAM, thus set bigger dimensions.") - 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.") - self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.") - self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 1.0", help_message="Forces the neural network to learn small details of the face. Enable it only when the face is trained enough with lr_dropout(on) and 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, 1.0 ) + self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 1.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, 1.0 ) if self.options['gan_power'] != 0.0: gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 ) @@ -151,17 +147,19 @@ class AMPModel(ModelBase): d_dims = self.options['d_dims'] d_mask_dims = self.options['d_mask_dims'] inter_res = self.inter_res = resolution // 32 - + use_fp16 = self.options['use_fp16'] + conv_dtype = tf.float16 if use_fp16 else tf.float32 + class Downscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=5 ): - self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME') + self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME', dtype=conv_dtype) def forward(self, x): return tf.nn.leaky_relu(self.conv1(x), 0.1) class Upscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=3 ): - self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME') + self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) def forward(self, x): x = nn.depth_to_space(tf.nn.leaky_relu(self.conv1(x), 0.1), 2) @@ -169,8 +167,8 @@ class AMPModel(ModelBase): class ResidualBlock(nn.ModelBase): def on_build(self, ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME') - self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME') + self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) + self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) def forward(self, inp): x = self.conv1(inp) @@ -190,8 +188,9 @@ class AMPModel(ModelBase): self.res5 = ResidualBlock(e_dims*8) self.dense1 = nn.Dense( (( resolution//(2**5) )**2) * e_dims*8, ae_dims ) - def forward(self, inp): - x = inp + def forward(self, x): + if use_fp16: + x = tf.cast(x, tf.float16) x = self.down1(x) x = self.res1(x) x = self.down2(x) @@ -199,6 +198,8 @@ class AMPModel(ModelBase): x = self.down4(x) x = self.down5(x) x = self.res5(x) + if use_fp16: + x = tf.cast(x, tf.float32) x = nn.pixel_norm(nn.flatten(x), axes=-1) x = self.dense1(x) return x @@ -232,15 +233,16 @@ class AMPModel(ModelBase): self.upscalem2 = Upscale(d_mask_dims*8, d_mask_dims*4, kernel_size=3) self.upscalem3 = Upscale(d_mask_dims*4, d_mask_dims*2, kernel_size=3) self.upscalem4 = Upscale(d_mask_dims*2, d_mask_dims*1, kernel_size=3) - self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME') + self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype) - self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME') - self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME') - self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME') - self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME') + self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype) + self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) + self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) + self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) - def forward(self, inp): - z = inp + def forward(self, z): + if use_fp16: + z = tf.cast(z, tf.float16) x = self.upscale0(z) x = self.res0(x) @@ -262,6 +264,10 @@ class AMPModel(ModelBase): m = self.upscalem3(m) m = self.upscalem4(m) m = tf.nn.sigmoid(self.out_convm(m)) + + if use_fp16: + x = tf.cast(x, tf.float32) + m = tf.cast(m, tf.float32) return x, m self.face_type = {'f' : FaceType.FULL, @@ -339,12 +345,12 @@ class AMPModel(ModelBase): if self.is_training: if gan_power != 0: - self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="GAN") + self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=use_fp16, name="GAN") self.model_filename_list += [ [self.GAN, 'GAN.npy'] ] # Initialize optimizers lr=5e-5 - lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] else 1.0 #and not self.pretrain + lr_dropout = 0.3 clipnorm = 1.0 if self.options['clipgrad'] else 0.0 @@ -355,12 +361,12 @@ class AMPModel(ModelBase): self.trainable_weights = self.encoder.get_weights() + self.decoder.get_weights() self.src_dst_opt = nn.AdaBelief(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt') - self.src_dst_opt.initialize_variables (self.all_weights, vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu') + self.src_dst_opt.initialize_variables (self.all_weights, vars_on_cpu=optimizer_vars_on_cpu) self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ] if gan_power != 0: self.GAN_opt = nn.AdaBelief(lr=lr, lr_dropout=lr_dropout, 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.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)#+self.D_src_x2.get_weights() self.model_filename_list += [ (self.GAN_opt, 'GAN_opt.npy') ] if self.is_training: diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 147defe..76c185b 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -29,7 +29,8 @@ class SAEHDModel(ModelBase): yn_str = {True:'y',False:'n'} min_res = 64 max_res = 640 - + + default_usefp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) default_resolution = self.options['resolution'] = self.load_or_def_option('resolution', 128) default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'f') default_models_opt_on_gpu = self.options['models_opt_on_gpu'] = self.load_or_def_option('models_opt_on_gpu', True) @@ -68,7 +69,8 @@ class SAEHDModel(ModelBase): self.ask_random_src_flip() self.ask_random_dst_flip() self.ask_batch_size(suggest_batch_size) - + self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters.') + if self.is_first_run(): resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16 and 32 for -d archi.") resolution = np.clip ( (resolution // 16) * 16, min_res, max_res) @@ -260,7 +262,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape, name='target_dstm_em') # Initializing model classes - model_archi = nn.DeepFakeArchi(resolution, opts=archi_opts) + model_archi = nn.DeepFakeArchi(resolution, use_fp16=self.options['use_fp16'], opts=archi_opts) with tf.device (models_opt_device): if 'df' in archi_type: @@ -301,7 +303,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.is_training: if gan_power != 0: - self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="D_src") + self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src") self.model_filename_list += [ [self.D_src, 'GAN.npy'] ] # Initialize optimizers From 2edac3df8cd334b5695c5e066cbb77b2992f5098 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 15 Jul 2021 00:24:38 +0400 Subject: [PATCH 16/30] fix for clipgrad --- core/leras/optimizers/AdaBelief.py | 4 ++-- core/leras/optimizers/RMSprop.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/core/leras/optimizers/AdaBelief.py b/core/leras/optimizers/AdaBelief.py index aed4308..dd004f8 100644 --- a/core/leras/optimizers/AdaBelief.py +++ b/core/leras/optimizers/AdaBelief.py @@ -50,11 +50,11 @@ class AdaBelief(nn.OptimizerBase): updates = [] if self.clipnorm > 0.0: - norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars])) + norm = tf.sqrt( sum([tf.reduce_sum(tf.square(tf.cast(g, tf.float32))) for g,v in grads_vars])) updates += [ state_ops.assign_add( self.iterations, 1) ] for i, (g,v) in enumerate(grads_vars): if self.clipnorm > 0.0: - g = self.tf_clip_norm(g, self.clipnorm, norm) + g = self.tf_clip_norm(g, self.clipnorm, tf.cast(norm, g.dtype) ) ms = self.ms_dict[ v.name ] vs = self.vs_dict[ v.name ] diff --git a/core/leras/optimizers/RMSprop.py b/core/leras/optimizers/RMSprop.py index 345b2a7..d2eb605 100644 --- a/core/leras/optimizers/RMSprop.py +++ b/core/leras/optimizers/RMSprop.py @@ -47,11 +47,11 @@ class RMSprop(nn.OptimizerBase): updates = [] if self.clipnorm > 0.0: - norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars])) + norm = tf.sqrt( sum([tf.reduce_sum(tf.square(tf.cast(g, tf.float32))) for g,v in grads_vars])) updates += [ state_ops.assign_add( self.iterations, 1) ] for i, (g,v) in enumerate(grads_vars): if self.clipnorm > 0.0: - g = self.tf_clip_norm(g, self.clipnorm, norm) + g = self.tf_clip_norm(g, self.clipnorm, tf.cast(norm, g.dtype) ) a = self.accumulators_dict[ v.name ] From f044c99ddc158cd23d090826f22957bdffd0ee06 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 15 Jul 2021 00:50:11 +0400 Subject: [PATCH 17/30] remove ConvolutionAwareInitializer from layers, because lr-dropout does the same thing. --- core/leras/layers/Conv2D.py | 4 ++-- core/leras/layers/Conv2DTranspose.py | 4 ++-- core/leras/layers/DepthwiseConv2D.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/core/leras/layers/Conv2D.py b/core/leras/layers/Conv2D.py index 7d4d444..a5febf0 100644 --- a/core/leras/layers/Conv2D.py +++ b/core/leras/layers/Conv2D.py @@ -55,8 +55,8 @@ class Conv2D(nn.LayerBase): if kernel_initializer is None: kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype) - if kernel_initializer is None: - kernel_initializer = nn.initializers.ca() + #if kernel_initializer is None: + # kernel_initializer = nn.initializers.ca() self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.out_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable ) diff --git a/core/leras/layers/Conv2DTranspose.py b/core/leras/layers/Conv2DTranspose.py index 937d624..a2e97dc 100644 --- a/core/leras/layers/Conv2DTranspose.py +++ b/core/leras/layers/Conv2DTranspose.py @@ -38,8 +38,8 @@ class Conv2DTranspose(nn.LayerBase): if kernel_initializer is None: kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype) - if kernel_initializer is None: - kernel_initializer = nn.initializers.ca() + #if kernel_initializer is None: + # kernel_initializer = nn.initializers.ca() self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.out_ch,self.in_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable ) if self.use_bias: diff --git a/core/leras/layers/DepthwiseConv2D.py b/core/leras/layers/DepthwiseConv2D.py index 98f3a0a..2916f01 100644 --- a/core/leras/layers/DepthwiseConv2D.py +++ b/core/leras/layers/DepthwiseConv2D.py @@ -68,8 +68,8 @@ class DepthwiseConv2D(nn.LayerBase): if kernel_initializer is None: kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype) - if kernel_initializer is None: - kernel_initializer = nn.initializers.ca() + #if kernel_initializer is None: + # kernel_initializer = nn.initializers.ca() self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.depth_multiplier), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable ) From 7b8991a3ccb03294f1246936ff68369bfc7676f4 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 15 Jul 2021 00:50:34 +0400 Subject: [PATCH 18/30] fix help --- models/Model_AMP/Model.py | 4 ++-- models/Model_SAEHD/Model.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index f6129ad..37aa828 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -65,7 +65,7 @@ class AMPModel(ModelBase): self.ask_random_src_flip() self.ask_random_dst_flip() self.ask_batch_size(suggest_batch_size) - self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters.') + self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters. Lr-dropout should be enabled.') if self.is_first_run(): resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 32 .") @@ -113,7 +113,7 @@ class AMPModel(ModelBase): self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.") - self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 1.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, 1.0 ) + 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 ) if self.options['gan_power'] != 0.0: gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 ) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 76c185b..820bba0 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -152,7 +152,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.") - self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 1.0", help_message="Forces the neural network to learn small details of the face. Enable it only when the face is trained enough with lr_dropout(on) and 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, 1.0 ) + 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 lr_dropout(on) and 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 ) if self.options['gan_power'] != 0.0: gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 ) From c9e0ba1779f1582973c9f9570a9eefab8ea65471 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 15 Jul 2021 13:22:32 +0400 Subject: [PATCH 19/30] removing fp16 for sae --- models/Model_SAEHD/Model.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 820bba0..4dcecd9 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -30,7 +30,7 @@ class SAEHDModel(ModelBase): min_res = 64 max_res = 640 - default_usefp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) + #default_usefp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) default_resolution = self.options['resolution'] = self.load_or_def_option('resolution', 128) default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'f') default_models_opt_on_gpu = self.options['models_opt_on_gpu'] = self.load_or_def_option('models_opt_on_gpu', True) @@ -69,7 +69,7 @@ class SAEHDModel(ModelBase): self.ask_random_src_flip() self.ask_random_dst_flip() self.ask_batch_size(suggest_batch_size) - self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters.') + #self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters.') if self.is_first_run(): resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16 and 32 for -d archi.") @@ -219,7 +219,8 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.set_iter(0) adabelief = self.options['adabelief'] - + use_fp16 = False#self.options['use_fp16'] + self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power'] random_warp = False if self.pretrain else self.options['random_warp'] random_src_flip = self.random_src_flip if not self.pretrain else True @@ -262,7 +263,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape, name='target_dstm_em') # Initializing model classes - model_archi = nn.DeepFakeArchi(resolution, use_fp16=self.options['use_fp16'], opts=archi_opts) + model_archi = nn.DeepFakeArchi(resolution, use_fp16=use_fp16, opts=archi_opts) with tf.device (models_opt_device): if 'df' in archi_type: @@ -303,7 +304,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.is_training: if gan_power != 0: - self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=self.options['use_fp16'], name="D_src") + self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="D_src") self.model_filename_list += [ [self.D_src, 'GAN.npy'] ] # Initialize optimizers From da8f33ee851eae62fa88113c3e1bb162a814f4c0 Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 15 Jul 2021 13:23:33 +0400 Subject: [PATCH 20/30] removing fp16 for amp --- models/Model_AMP/Model.py | 38 ++++++++++++++++++++------------------ 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 37aa828..fc3f28f 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -31,7 +31,7 @@ class AMPModel(ModelBase): min_res = 64 max_res = 640 - default_usefp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) + #default_usefp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) default_resolution = self.options['resolution'] = self.load_or_def_option('resolution', 224) default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'wf') default_models_opt_on_gpu = self.options['models_opt_on_gpu'] = self.load_or_def_option('models_opt_on_gpu', True) @@ -65,7 +65,7 @@ class AMPModel(ModelBase): self.ask_random_src_flip() self.ask_random_dst_flip() self.ask_batch_size(suggest_batch_size) - self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters. Lr-dropout should be enabled.') + #self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters. Lr-dropout should be enabled.') if self.is_first_run(): resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 32 .") @@ -147,19 +147,21 @@ class AMPModel(ModelBase): d_dims = self.options['d_dims'] d_mask_dims = self.options['d_mask_dims'] inter_res = self.inter_res = resolution // 32 - use_fp16 = self.options['use_fp16'] - conv_dtype = tf.float16 if use_fp16 else tf.float32 + use_fp16 = False# self.options['use_fp16'] + + ae_use_fp16 = use_fp16 + ae_conv_dtype = tf.float16 if use_fp16 else tf.float32 class Downscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=5 ): - self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME', dtype=conv_dtype) + self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME', dtype=ae_conv_dtype) def forward(self, x): return tf.nn.leaky_relu(self.conv1(x), 0.1) class Upscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=3 ): - self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) + self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=ae_conv_dtype) def forward(self, x): x = nn.depth_to_space(tf.nn.leaky_relu(self.conv1(x), 0.1), 2) @@ -167,8 +169,8 @@ class AMPModel(ModelBase): class ResidualBlock(nn.ModelBase): def on_build(self, ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) - self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) + self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=ae_conv_dtype) + self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=ae_conv_dtype) def forward(self, inp): x = self.conv1(inp) @@ -189,7 +191,7 @@ class AMPModel(ModelBase): self.dense1 = nn.Dense( (( resolution//(2**5) )**2) * e_dims*8, ae_dims ) def forward(self, x): - if use_fp16: + if ae_use_fp16: x = tf.cast(x, tf.float16) x = self.down1(x) x = self.res1(x) @@ -198,7 +200,7 @@ class AMPModel(ModelBase): x = self.down4(x) x = self.down5(x) x = self.res5(x) - if use_fp16: + if ae_use_fp16: x = tf.cast(x, tf.float32) x = nn.pixel_norm(nn.flatten(x), axes=-1) x = self.dense1(x) @@ -233,15 +235,15 @@ class AMPModel(ModelBase): self.upscalem2 = Upscale(d_mask_dims*8, d_mask_dims*4, kernel_size=3) self.upscalem3 = Upscale(d_mask_dims*4, d_mask_dims*2, kernel_size=3) self.upscalem4 = Upscale(d_mask_dims*2, d_mask_dims*1, kernel_size=3) - self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype) + self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME', dtype=ae_conv_dtype) - self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype) - self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) - self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) - self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) + self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME', dtype=ae_conv_dtype) + self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=ae_conv_dtype) + self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=ae_conv_dtype) + self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=ae_conv_dtype) def forward(self, z): - if use_fp16: + if ae_use_fp16: z = tf.cast(z, tf.float16) x = self.upscale0(z) @@ -265,7 +267,7 @@ class AMPModel(ModelBase): m = self.upscalem4(m) m = tf.nn.sigmoid(self.out_convm(m)) - if use_fp16: + if ae_use_fp16: x = tf.cast(x, tf.float32) m = tf.cast(m, tf.float32) return x, m @@ -345,7 +347,7 @@ class AMPModel(ModelBase): if self.is_training: if gan_power != 0: - self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], use_fp16=use_fp16, name="GAN") + self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="GAN") self.model_filename_list += [ [self.GAN, 'GAN.npy'] ] # Initialize optimizers From 959a3530f8606e678fd087a4471bb7b8daf118e2 Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 17 Jul 2021 10:27:07 +0400 Subject: [PATCH 21/30] AMP: default morph factor is now 0.5. Preview shows with 1.0 --- models/Model_AMP/Model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index fc3f28f..37789a3 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -46,7 +46,7 @@ class AMPModel(ModelBase): default_e_dims = self.options['e_dims'] = self.load_or_def_option('e_dims', 64) default_d_dims = self.options['d_dims'] = self.options.get('d_dims', None) default_d_mask_dims = self.options['d_mask_dims'] = self.options.get('d_mask_dims', None) - default_morph_factor = self.options['morph_factor'] = self.options.get('morph_factor', 0.1) + default_morph_factor = self.options['morph_factor'] = self.options.get('morph_factor', 0.5) default_masked_training = self.options['masked_training'] = self.load_or_def_option('masked_training', True) default_eyes_mouth_prio = self.options['eyes_mouth_prio'] = self.load_or_def_option('eyes_mouth_prio', True) default_uniform_yaw = self.options['uniform_yaw'] = self.load_or_def_option('uniform_yaw', False) @@ -780,9 +780,9 @@ class AMPModel(ModelBase): i = np.random.randint(n_samples) if not for_history else 0 st = [ np.concatenate ((S[i], D[i], DD[i]*DDM_000[i]), axis=1) ] - st += [ np.concatenate ((SS[i], DD[i], SD_075[i] ), axis=1) ] + st += [ np.concatenate ((SS[i], DD[i], SD_100[i] ), axis=1) ] - result += [ ('AMP morph 0.75', np.concatenate (st, axis=0 )), ] + result += [ ('AMP morph 1.0', np.concatenate (st, axis=0 )), ] 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) ] @@ -803,7 +803,7 @@ class AMPModel(ModelBase): #override def get_MergerConfig(self): - morph_factor = np.clip ( io.input_number ("Morph factor", 0.75, add_info="0.0 .. 1.0"), 0.0, 1.0 ) + morph_factor = np.clip ( io.input_number ("Morph factor", 1.0, add_info="0.0 .. 1.0"), 0.0, 1.0 ) def predictor_morph(face): return self.predictor_func(face, morph_factor) From 0748b8d04351c65d4c6b49e5a9bd24258f476ce3 Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 17 Jul 2021 21:58:57 +0400 Subject: [PATCH 22/30] AMP: removed eyes mouth prio option(default enabled), removed masked training option(default enabled). --- models/Model_AMP/Model.py | 372 ++++++++++++++------------------------ 1 file changed, 137 insertions(+), 235 deletions(-) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 37789a3..0a5bc18 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -18,44 +18,26 @@ class AMPModel(ModelBase): def on_initialize_options(self): device_config = nn.getCurrentDeviceConfig() - lowest_vram = 2 - if len(device_config.devices) != 0: - lowest_vram = device_config.devices.get_worst_device().total_mem_gb - - if lowest_vram >= 4: - suggest_batch_size = 8 - else: - suggest_batch_size = 4 - - yn_str = {True:'y',False:'n'} - min_res = 64 - max_res = 640 - - #default_usefp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) default_resolution = self.options['resolution'] = self.load_or_def_option('resolution', 224) default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'wf') default_models_opt_on_gpu = self.options['models_opt_on_gpu'] = self.load_or_def_option('models_opt_on_gpu', True) default_ae_dims = self.options['ae_dims'] = self.load_or_def_option('ae_dims', 256) - + inter_dims = self.load_or_def_option('inter_dims', None) if inter_dims is None: inter_dims = self.options['ae_dims'] - default_inter_dims = self.options['inter_dims'] = inter_dims - + default_inter_dims = self.options['inter_dims'] = inter_dims + default_e_dims = self.options['e_dims'] = self.load_or_def_option('e_dims', 64) default_d_dims = self.options['d_dims'] = self.options.get('d_dims', None) default_d_mask_dims = self.options['d_mask_dims'] = self.options.get('d_mask_dims', None) default_morph_factor = self.options['morph_factor'] = self.options.get('morph_factor', 0.5) - default_masked_training = self.options['masked_training'] = self.load_or_def_option('masked_training', True) - default_eyes_mouth_prio = self.options['eyes_mouth_prio'] = self.load_or_def_option('eyes_mouth_prio', True) default_uniform_yaw = self.options['uniform_yaw'] = self.load_or_def_option('uniform_yaw', False) default_random_warp = self.options['random_warp'] = self.load_or_def_option('random_warp', True) default_ct_mode = self.options['ct_mode'] = self.load_or_def_option('ct_mode', 'none') default_clipgrad = self.options['clipgrad'] = self.load_or_def_option('clipgrad', False) - #default_pretrain = self.options['pretrain'] = self.load_or_def_option('pretrain', False) - ask_override = self.ask_override() if self.is_first_run() or ask_override: @@ -64,12 +46,11 @@ class AMPModel(ModelBase): self.ask_target_iter() self.ask_random_src_flip() self.ask_random_dst_flip() - self.ask_batch_size(suggest_batch_size) - #self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters. Lr-dropout should be enabled.') - + self.ask_batch_size(8) + if self.is_first_run(): resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 32 .") - resolution = np.clip ( (resolution // 32) * 32, min_res, max_res) + resolution = np.clip ( (resolution // 32) * 32, 64, 640) self.options['resolution'] = resolution self.options['face_type'] = io.input_str ("Face type", default_face_type, ['f','wf','head'], help_message="whole face / head").lower() @@ -93,15 +74,10 @@ class AMPModel(ModelBase): d_mask_dims = np.clip ( io.input_int("Decoder mask dimensions", default_d_mask_dims, add_info="16-256", help_message="Typical mask dimensions = decoder dimensions / 3. If you manually cut out obstacles from the dst mask, you can increase this parameter to achieve better quality." ), 16, 256 ) self.options['d_mask_dims'] = d_mask_dims + d_mask_dims % 2 - morph_factor = np.clip ( io.input_number ("Morph factor.", default_morph_factor, add_info="0.1 .. 0.5", help_message="The smaller the value, the more src-like facial expressions will appear. The larger the value, the less space there is to train a large dst faceset in the neural network. Typical fine value is 0.33"), 0.1, 0.5 ) + morph_factor = np.clip ( io.input_number ("Morph factor.", default_morph_factor, add_info="0.1 .. 0.5", help_message="Typical fine value is 0.5"), 0.1, 0.5 ) self.options['morph_factor'] = morph_factor - if self.is_first_run() or ask_override: - if self.options['face_type'] == 'wf' or self.options['face_type'] == 'head': - self.options['masked_training'] = io.input_bool ("Masked training", default_masked_training, help_message="This option is available only for 'whole_face' or 'head' type. Masked training clips training area to full_face mask or XSeg mask, thus network will train the faces properly.") - - self.options['eyes_mouth_prio'] = io.input_bool ("Eyes and mouth priority", default_eyes_mouth_prio, help_message='Helps to fix eye problems during training like "alien eyes" and wrong eyes direction. Also makes the detail of the teeth higher.') self.options['uniform_yaw'] = io.input_bool ("Uniform yaw distribution of samples", default_uniform_yaw, help_message='Helps to fix blurry side faces due to small amount of them in the faceset.') default_gan_power = self.options['gan_power'] = self.load_or_def_option('gan_power', 0.0) @@ -125,10 +101,7 @@ class AMPModel(ModelBase): self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") self.options['clipgrad'] = io.input_bool ("Enable gradient clipping", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.") - #self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly. Forces random_warp=N, random_flips=Y, gan_power=0.0, lr_dropout=N, uniform_yaw=Y") - self.gan_model_changed = (default_gan_patch_size != self.options['gan_patch_size']) or (default_gan_dims != self.options['gan_dims']) - #self.pretrain_just_disabled = (default_pretrain == True and self.options['pretrain'] == False) #override def on_initialize(self): @@ -138,30 +111,38 @@ class AMPModel(ModelBase): nn.initialize(data_format=self.model_data_format) tf = nn.tf - self.resolution = resolution = self.options['resolution'] - input_ch=3 - ae_dims = self.ae_dims = self.options['ae_dims'] - inter_dims = self.inter_dims = self.options['inter_dims'] - e_dims = self.options['e_dims'] - d_dims = self.options['d_dims'] + resolution = self.resolution = self.options['resolution'] + e_dims = self.options['e_dims'] + ae_dims = self.options['ae_dims'] + inter_dims = self.inter_dims = self.options['inter_dims'] + inter_res = self.inter_res = resolution // 32 + d_dims = self.options['d_dims'] d_mask_dims = self.options['d_mask_dims'] - inter_res = self.inter_res = resolution // 32 - use_fp16 = False# self.options['use_fp16'] - - ae_use_fp16 = use_fp16 - ae_conv_dtype = tf.float16 if use_fp16 else tf.float32 - + face_type = self.face_type = {'f' : FaceType.FULL, + 'wf' : FaceType.WHOLE_FACE, + 'head' : FaceType.HEAD}[ self.options['face_type'] ] + morph_factor = self.options['morph_factor'] + gan_power = self.gan_power = self.options['gan_power'] + random_warp = self.options['random_warp'] + + ct_mode = self.options['ct_mode'] + if ct_mode == 'none': + ct_mode = None + + use_fp16 = self.is_exporting + conv_dtype = tf.float16 if use_fp16 else tf.float32 + class Downscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=5 ): - self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME', dtype=ae_conv_dtype) + self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME', dtype=conv_dtype) def forward(self, x): return tf.nn.leaky_relu(self.conv1(x), 0.1) class Upscale(nn.ModelBase): def on_build(self, in_ch, out_ch, kernel_size=3 ): - self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=ae_conv_dtype) + self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) def forward(self, x): x = nn.depth_to_space(tf.nn.leaky_relu(self.conv1(x), 0.1), 2) @@ -169,8 +150,8 @@ class AMPModel(ModelBase): class ResidualBlock(nn.ModelBase): def on_build(self, ch, kernel_size=3 ): - self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=ae_conv_dtype) - self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=ae_conv_dtype) + self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) + self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype) def forward(self, inp): x = self.conv1(inp) @@ -191,7 +172,7 @@ class AMPModel(ModelBase): self.dense1 = nn.Dense( (( resolution//(2**5) )**2) * e_dims*8, ae_dims ) def forward(self, x): - if ae_use_fp16: + if use_fp16: x = tf.cast(x, tf.float16) x = self.down1(x) x = self.res1(x) @@ -200,7 +181,7 @@ class AMPModel(ModelBase): x = self.down4(x) x = self.down5(x) x = self.res5(x) - if ae_use_fp16: + if use_fp16: x = tf.cast(x, tf.float32) x = nn.pixel_norm(nn.flatten(x), axes=-1) x = self.dense1(x) @@ -235,15 +216,15 @@ class AMPModel(ModelBase): self.upscalem2 = Upscale(d_mask_dims*8, d_mask_dims*4, kernel_size=3) self.upscalem3 = Upscale(d_mask_dims*4, d_mask_dims*2, kernel_size=3) self.upscalem4 = Upscale(d_mask_dims*2, d_mask_dims*1, kernel_size=3) - self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME', dtype=ae_conv_dtype) + self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype) - self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME', dtype=ae_conv_dtype) - self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=ae_conv_dtype) - self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=ae_conv_dtype) - self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=ae_conv_dtype) + self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype) + self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) + self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) + self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype) def forward(self, z): - if ae_use_fp16: + if use_fp16: z = tf.cast(z, tf.float16) x = self.upscale0(z) @@ -259,60 +240,22 @@ class AMPModel(ModelBase): self.out_conv1(x), self.out_conv2(x), self.out_conv3(x)), nn.conv2d_ch_axis), 2) ) - m = self.upscalem0(z) m = self.upscalem1(m) m = self.upscalem2(m) m = self.upscalem3(m) m = self.upscalem4(m) m = tf.nn.sigmoid(self.out_convm(m)) - - if ae_use_fp16: + + if use_fp16: x = tf.cast(x, tf.float32) m = tf.cast(m, tf.float32) return x, m - self.face_type = {'f' : FaceType.FULL, - 'wf' : FaceType.WHOLE_FACE, - 'head' : FaceType.HEAD}[ self.options['face_type'] ] - - if 'eyes_prio' in self.options: - self.options.pop('eyes_prio') - - eyes_mouth_prio = self.options['eyes_mouth_prio'] - - - morph_factor = self.options['morph_factor'] - - gan_power = self.gan_power = self.options['gan_power'] - random_warp = self.options['random_warp'] - random_src_flip = self.random_src_flip - random_dst_flip = self.random_dst_flip - - #pretrain = self.pretrain = self.options['pretrain'] - #if self.pretrain_just_disabled: - # self.set_iter(0) - # self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power'] - # random_warp = False if self.pretrain else self.options['random_warp'] - # random_src_flip = self.random_src_flip if not self.pretrain else True - # random_dst_flip = self.random_dst_flip if not self.pretrain else True - - # if self.pretrain: - # self.options_show_override['gan_power'] = 0.0 - # self.options_show_override['random_warp'] = False - # self.options_show_override['lr_dropout'] = 'n' - # self.options_show_override['uniform_yaw'] = True - - masked_training = self.options['masked_training'] - ct_mode = self.options['ct_mode'] - if ct_mode == 'none': - ct_mode = None - models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu'] models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0' optimizer_vars_on_cpu = models_opt_device=='/CPU:0' - bgr_shape = self.bgr_shape = nn.get4Dshape(resolution,resolution,input_ch) mask_shape = nn.get4Dshape(resolution,resolution,1) self.model_filename_list = [] @@ -333,7 +276,6 @@ class AMPModel(ModelBase): self.morph_value_t = tf.placeholder (nn.floatx, (1,), name='morph_value_t') # Initializing model classes - with tf.device (models_opt_device): self.encoder = Encoder(name='encoder') self.inter_src = Inter(name='inter_src') @@ -346,30 +288,21 @@ class AMPModel(ModelBase): [self.decoder , 'decoder.npy'] ] if self.is_training: - if gan_power != 0: - self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="GAN") - self.model_filename_list += [ [self.GAN, 'GAN.npy'] ] - # Initialize optimizers - lr=5e-5 - lr_dropout = 0.3 - clipnorm = 1.0 if self.options['clipgrad'] else 0.0 self.all_weights = self.encoder.get_weights() + self.decoder.get_weights() - #if pretrain: - # self.trainable_weights = self.encoder.get_weights() + self.decoder.get_weights() - #else: - self.trainable_weights = self.encoder.get_weights() + self.decoder.get_weights() - self.src_dst_opt = nn.AdaBelief(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt') + self.src_dst_opt = nn.AdaBelief(lr=5e-5, lr_dropout=0.3, clipnorm=clipnorm, name='src_dst_opt') self.src_dst_opt.initialize_variables (self.all_weights, vars_on_cpu=optimizer_vars_on_cpu) self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ] if gan_power != 0: - self.GAN_opt = nn.AdaBelief(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='GAN_opt') - self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)#+self.D_src_x2.get_weights() - self.model_filename_list += [ (self.GAN_opt, 'GAN_opt.npy') ] + self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="GAN") + self.GAN_opt = nn.AdaBelief(lr=5e-5, lr_dropout=0.3, clipnorm=clipnorm, name='GAN_opt') + self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu) + self.model_filename_list += [ [self.GAN, 'GAN.npy'], + [self.GAN_opt, 'GAN_opt.npy'] ] if self.is_training: # Adjust batch size for multiple GPU @@ -387,8 +320,8 @@ class AMPModel(ModelBase): gpu_src_losses = [] gpu_dst_losses = [] - gpu_G_loss_gvs = [] - gpu_D_src_dst_loss_gvs = [] + gpu_G_loss_gradients = [] + gpu_GAN_loss_grads = [] for gpu_id in range(gpu_count): with tf.device( f'/{devices[gpu_id].tf_dev_type}:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ): @@ -408,16 +341,8 @@ class AMPModel(ModelBase): gpu_src_code = self.encoder (gpu_warped_src) gpu_dst_code = self.encoder (gpu_warped_dst) - # if pretrain: - # gpu_src_inter_src_code = self.inter_src (gpu_src_code) - # gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) - # gpu_src_code = gpu_src_inter_src_code * nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor) - # gpu_dst_code = gpu_src_dst_code = gpu_dst_inter_dst_code * nn.random_binomial( [bs_per_gpu, gpu_dst_inter_dst_code.shape.as_list()[1], 1,1] , p=morph_factor) - # else: - gpu_src_inter_src_code = self.inter_src (gpu_src_code) - gpu_src_inter_dst_code = self.inter_dst (gpu_src_code) - gpu_dst_inter_src_code = self.inter_src (gpu_dst_code) - gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) + gpu_src_inter_src_code, gpu_src_inter_dst_code = self.inter_src (gpu_src_code), self.inter_dst (gpu_src_code) + gpu_dst_inter_src_code, gpu_dst_inter_dst_code = self.inter_src (gpu_dst_code), self.inter_dst (gpu_dst_code) inter_rnd_binomial = nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor) gpu_src_code = gpu_src_inter_src_code * inter_rnd_binomial + gpu_src_inter_dst_code * (1-inter_rnd_binomial) @@ -431,62 +356,50 @@ class AMPModel(ModelBase): gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code) - gpu_pred_src_src_list.append(gpu_pred_src_src) - gpu_pred_dst_dst_list.append(gpu_pred_dst_dst) - gpu_pred_src_dst_list.append(gpu_pred_src_dst) + gpu_pred_src_src_list.append(gpu_pred_src_src), gpu_pred_src_srcm_list.append(gpu_pred_src_srcm) + gpu_pred_dst_dst_list.append(gpu_pred_dst_dst), gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm) + gpu_pred_src_dst_list.append(gpu_pred_src_dst), gpu_pred_src_dstm_list.append(gpu_pred_src_dstm) - gpu_pred_src_srcm_list.append(gpu_pred_src_srcm) - gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm) - gpu_pred_src_dstm_list.append(gpu_pred_src_dstm) + gpu_target_srcm_blur = tf.clip_by_value( nn.gaussian_blur(gpu_target_srcm, max(1, resolution // 32) ), 0, 0.5) * 2 + gpu_target_dstm_blur = tf.clip_by_value(nn.gaussian_blur(gpu_target_dstm, max(1, resolution // 32) ), 0, 0.5) * 2 - gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm, max(1, resolution // 32) ) - gpu_target_srcm_blur = tf.clip_by_value(gpu_target_srcm_blur, 0, 0.5) * 2 + gpu_target_srcm_anti_blur = 1.0-gpu_target_srcm_blur + gpu_target_dstm_anti_blur = 1.0-gpu_target_dstm_blur - gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm, max(1, resolution // 32) ) - gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_blur, 0, 0.5) * 2 + gpu_target_src_masked = gpu_target_src*gpu_target_srcm_blur + gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur + gpu_target_src_anti_masked = gpu_target_src*gpu_target_srcm_anti_blur + gpu_target_dst_anti_masked = gpu_target_dst*gpu_target_dstm_anti_blur - gpu_target_dst_anti_masked = gpu_target_dst*(1.0-gpu_target_dstm_blur) - gpu_target_src_anti_masked = gpu_target_src*(1.0-gpu_target_srcm_blur) - gpu_target_src_masked_opt = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src - gpu_target_dst_masked_opt = gpu_target_dst*gpu_target_dstm_blur if masked_training else gpu_target_dst + gpu_pred_src_src_masked = gpu_pred_src_src*gpu_target_srcm_blur + gpu_pred_dst_dst_masked = gpu_pred_dst_dst*gpu_target_dstm_blur + gpu_pred_src_src_anti_masked = gpu_pred_src_src*gpu_target_srcm_anti_blur + gpu_pred_dst_dst_anti_masked = gpu_pred_dst_dst*gpu_target_dstm_anti_blur - gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src - gpu_pred_src_src_anti_masked = gpu_pred_src_src*(1.0-gpu_target_srcm_blur) - gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst - gpu_pred_dst_dst_anti_masked = gpu_pred_dst_dst*(1.0-gpu_target_dstm_blur) + # Structural loss + gpu_src_loss = tf.reduce_mean (5*nn.dssim(gpu_target_src_masked, gpu_pred_src_src_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) + gpu_src_loss += tf.reduce_mean (5*nn.dssim(gpu_target_src_masked, gpu_pred_src_src_masked, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) + gpu_dst_loss = tf.reduce_mean (5*nn.dssim(gpu_target_dst_masked, gpu_pred_dst_dst_masked, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) + gpu_dst_loss += tf.reduce_mean (5*nn.dssim(gpu_target_dst_masked, gpu_pred_dst_dst_masked, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1]) - if resolution < 256: - gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) - else: - gpu_dst_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1]) - gpu_dst_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1]) - gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3]) - if eyes_mouth_prio: - gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_dstm_em - gpu_pred_dst_dst*gpu_target_dstm_em ), axis=[1,2,3]) - gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) - gpu_dst_loss += 0.1*tf.reduce_mean(tf.square(gpu_pred_dst_dst_anti_masked-gpu_target_dst_anti_masked),axis=[1,2,3] ) - gpu_dst_losses += [gpu_dst_loss] + # Pixel loss + gpu_src_loss += tf.reduce_mean (10*tf.square(gpu_target_src_masked-gpu_pred_src_src_masked), axis=[1,2,3]) + gpu_dst_loss += tf.reduce_mean (10*tf.square(gpu_target_dst_masked-gpu_pred_dst_dst_masked), axis=[1,2,3]) - #if not pretrain: - if resolution < 256: - gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) - else: - gpu_src_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1]) - gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1]) - gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3]) - - if eyes_mouth_prio: - gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_em - gpu_pred_src_src*gpu_target_srcm_em ), axis=[1,2,3]) + # Eyes+mouth prio loss + gpu_src_loss += tf.reduce_mean (300*tf.abs (gpu_target_src*gpu_target_srcm_em-gpu_pred_src_src*gpu_target_srcm_em), axis=[1,2,3]) + gpu_dst_loss += tf.reduce_mean (300*tf.abs (gpu_target_dst*gpu_target_dstm_em-gpu_pred_dst_dst*gpu_target_dstm_em), axis=[1,2,3]) + # Mask loss gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] ) - #else: - # gpu_src_loss = gpu_dst_loss + gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] ) + + # dst-dst background weak loss + gpu_dst_loss += tf.reduce_mean(0.1*tf.square(gpu_pred_dst_dst_anti_masked-gpu_target_dst_anti_masked),axis=[1,2,3] ) + gpu_dst_loss += 0.000001*nn.total_variation_mse(gpu_pred_dst_dst_anti_masked) gpu_src_losses += [gpu_src_loss] - - #if pretrain: - # gpu_G_loss = gpu_dst_loss - #else: + gpu_dst_losses += [gpu_dst_loss] gpu_G_loss = gpu_src_loss + gpu_dst_loss def DLossOnes(logits): @@ -496,30 +409,28 @@ class AMPModel(ModelBase): return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(logits), logits=logits), axis=[1,2,3]) if gan_power != 0: - gpu_pred_src_src_d, gpu_pred_src_src_d2 = self.GAN(gpu_pred_src_src_masked_opt) - gpu_pred_dst_dst_d, gpu_pred_dst_dst_d2 = self.GAN(gpu_pred_dst_dst_masked_opt) - gpu_target_src_d, gpu_target_src_d2 = self.GAN(gpu_target_src_masked_opt) - gpu_target_dst_d, gpu_target_dst_d2 = self.GAN(gpu_target_dst_masked_opt) + gpu_pred_src_src_d, gpu_pred_src_src_d2 = self.GAN(gpu_pred_src_src_masked) + gpu_pred_dst_dst_d, gpu_pred_dst_dst_d2 = self.GAN(gpu_pred_dst_dst_masked) + gpu_target_src_d, gpu_target_src_d2 = self.GAN(gpu_target_src_masked) + gpu_target_dst_d, gpu_target_dst_d2 = self.GAN(gpu_target_dst_masked) - gpu_D_src_dst_loss = (DLossOnes (gpu_target_src_d) + DLossOnes (gpu_target_src_d2) + \ - DLossZeros(gpu_pred_src_src_d) + DLossZeros(gpu_pred_src_src_d2) + \ - DLossOnes (gpu_target_dst_d) + DLossOnes (gpu_target_dst_d2) + \ - DLossZeros(gpu_pred_dst_dst_d) + DLossZeros(gpu_pred_dst_dst_d2) - ) * ( 1.0 / 8) + gpu_GAN_loss = (DLossOnes (gpu_target_src_d) + DLossOnes (gpu_target_src_d2) + \ + DLossZeros(gpu_pred_src_src_d) + DLossZeros(gpu_pred_src_src_d2) + \ + DLossOnes (gpu_target_dst_d) + DLossOnes (gpu_target_dst_d2) + \ + DLossZeros(gpu_pred_dst_dst_d) + DLossZeros(gpu_pred_dst_dst_d2) + ) * (1.0 / 8) - gpu_D_src_dst_loss_gvs += [ nn.gradients (gpu_D_src_dst_loss, self.GAN.get_weights() ) ] + gpu_GAN_loss_grads += [ nn.gradients (gpu_GAN_loss, self.GAN.get_weights() ) ] gpu_G_loss += (DLossOnes(gpu_pred_src_src_d) + DLossOnes(gpu_pred_src_src_d2) + \ DLossOnes(gpu_pred_dst_dst_d) + DLossOnes(gpu_pred_dst_dst_d2) ) * gan_power - if masked_training: - # Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan - gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_src_src) - gpu_G_loss += 0.02*tf.reduce_mean(tf.square(gpu_pred_src_src_anti_masked-gpu_target_src_anti_masked),axis=[1,2,3] ) - - gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.trainable_weights ) ] + # Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan + gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_src_src) + gpu_G_loss += 0.02*tf.reduce_mean(tf.square(gpu_pred_src_src_anti_masked-gpu_target_src_anti_masked),axis=[1,2,3] ) + gpu_G_loss_gradients += [ nn.gradients ( gpu_G_loss, self.encoder.get_weights() + self.decoder.get_weights() ) ] # Average losses and gradients, and create optimizer update ops with tf.device(f'/CPU:0'): @@ -533,15 +444,15 @@ class AMPModel(ModelBase): with tf.device (models_opt_device): src_loss = tf.concat(gpu_src_losses, 0) dst_loss = tf.concat(gpu_dst_losses, 0) - src_dst_loss_gv_op = self.src_dst_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gvs)) + train_op = self.src_dst_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gradients)) if gan_power != 0: - src_D_src_dst_loss_gv_op = self.GAN_opt.get_update_op (nn.average_gv_list(gpu_D_src_dst_loss_gvs) ) + GAN_train_op = self.GAN_opt.get_update_op (nn.average_gv_list(gpu_GAN_loss_grads) ) # Initializing training and view functions - def src_dst_train(warped_src, target_src, target_srcm, target_srcm_em, \ + def train(warped_src, target_src, target_srcm, target_srcm_em, \ warped_dst, target_dst, target_dstm, target_dstm_em, ): - s, d, _ = nn.tf_sess.run ( [ src_loss, dst_loss, src_dst_loss_gv_op], + s, d, _ = nn.tf_sess.run ([src_loss, dst_loss, train_op], feed_dict={self.warped_src :warped_src, self.target_src :target_src, self.target_srcm:target_srcm, @@ -552,20 +463,20 @@ class AMPModel(ModelBase): self.target_dstm_em:target_dstm_em, }) return s, d - self.src_dst_train = src_dst_train + self.train = train if gan_power != 0: - def D_src_dst_train(warped_src, target_src, target_srcm, target_srcm_em, \ - warped_dst, target_dst, target_dstm, target_dstm_em, ): - nn.tf_sess.run ([src_D_src_dst_loss_gv_op], feed_dict={self.warped_src :warped_src, - self.target_src :target_src, - self.target_srcm:target_srcm, - self.target_srcm_em:target_srcm_em, - self.warped_dst :warped_dst, - self.target_dst :target_dst, - self.target_dstm:target_dstm, - self.target_dstm_em:target_dstm_em}) - self.D_src_dst_train = D_src_dst_train + def GAN_train(warped_src, target_src, target_srcm, target_srcm_em, \ + warped_dst, target_dst, target_dstm, target_dstm_em, ): + nn.tf_sess.run ([GAN_train_op], feed_dict={self.warped_src :warped_src, + self.target_src :target_src, + self.target_srcm:target_srcm, + self.target_srcm_em:target_srcm_em, + self.warped_dst :warped_dst, + self.target_dst :target_dst, + self.target_dstm:target_dstm, + self.target_dstm_em:target_dstm_em}) + self.GAN_train = GAN_train def AE_view(warped_src, warped_dst, morph_value): return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm], @@ -576,8 +487,8 @@ class AMPModel(ModelBase): #Initializing merge function with tf.device( nn.tf_default_device_name if len(devices) != 0 else f'/CPU:0'): gpu_dst_code = self.encoder (self.warped_dst) - gpu_dst_inter_src_code = self.inter_src ( gpu_dst_code) - gpu_dst_inter_dst_code = self.inter_dst ( gpu_dst_code) + gpu_dst_inter_src_code = self.inter_src (gpu_dst_code) + gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code) inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32) gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , inter_res, inter_res]), @@ -593,16 +504,10 @@ class AMPModel(ModelBase): # Loading/initializing all models/optimizers weights for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"): - # if self.pretrain_just_disabled: - # do_init = False - # if model == self.inter_src or model == self.inter_dst: - # do_init = True - # else: do_init = self.is_first_run() if self.is_training and gan_power != 0 and model == self.GAN: if self.gan_model_changed: do_init = True - if not do_init: do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) ) if do_init: @@ -614,7 +519,7 @@ class AMPModel(ModelBase): training_data_src_path = self.training_data_src_path #if not self.pretrain else self.get_pretraining_data_path() training_data_dst_path = self.training_data_dst_path #if not self.pretrain else self.get_pretraining_data_path() - random_ct_samples_path=training_data_dst_path if ct_mode is not None else None #and not self.pretrain + random_ct_samples_path=training_data_dst_path if ct_mode is not None else None #and not self.pretrain cpu_count = min(multiprocessing.cpu_count(), 8) src_generators_count = cpu_count // 2 @@ -624,21 +529,21 @@ class AMPModel(ModelBase): self.set_training_data_generators ([ SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(), - sample_process_options=SampleProcessor.Options(random_flip=random_src_flip), - output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, - {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, - {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, - {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, + sample_process_options=SampleProcessor.Options(random_flip=self.random_src_flip), + output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, + {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, + {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, + {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, ], uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain, generators_count=src_generators_count ), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(), - sample_process_options=SampleProcessor.Options(random_flip=random_dst_flip), - output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, - {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, - {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, - {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, + sample_process_options=SampleProcessor.Options(random_flip=self.random_dst_flip), + output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, + {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, + {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, + {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution}, ], uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain, generators_count=dst_generators_count ) @@ -646,15 +551,12 @@ class AMPModel(ModelBase): self.last_src_samples_loss = [] self.last_dst_samples_loss = [] - #if self.pretrain_just_disabled: - # self.update_sample_for_preview(force_new=True) - def export_dfm (self): output_path=self.get_strpath_storage_for_file('model.dfm') - + io.log_info(f'Dumping .dfm to {output_path}') - + tf = nn.tf with tf.device (nn.tf_default_device_name): warped_dst = tf.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face') @@ -679,13 +581,13 @@ class AMPModel(ModelBase): tf.identity(gpu_pred_dst_dstm, name='out_face_mask') tf.identity(gpu_pred_src_dst, name='out_celeb_face') tf.identity(gpu_pred_src_dstm, name='out_celeb_face_mask') - + output_graph_def = tf.graph_util.convert_variables_to_constants( - nn.tf_sess, - tf.get_default_graph().as_graph_def(), + nn.tf_sess, + tf.get_default_graph().as_graph_def(), ['out_face_mask','out_celeb_face','out_celeb_face_mask'] - ) - + ) + import tf2onnx with tf.device("/CPU:0"): model_proto, _ = tf2onnx.convert._convert_common( @@ -717,7 +619,7 @@ class AMPModel(ModelBase): ( (warped_src, target_src, target_srcm, target_srcm_em), \ (warped_dst, target_dst, target_dstm, target_dstm_em) ) = self.generate_next_samples() - src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) + src_loss, dst_loss = self.train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) for i in range(bs): self.last_src_samples_loss.append ( (src_loss[i], warped_src[i], target_src[i], target_srcm[i], target_srcm_em[i]) ) @@ -737,12 +639,12 @@ class AMPModel(ModelBase): target_dstm = np.stack( [ x[3] for x in dst_samples_loss[:bs] ] ) target_dstm_em = np.stack( [ x[4] for x in dst_samples_loss[:bs] ] ) - src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) + src_loss, dst_loss = self.train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) self.last_src_samples_loss = [] self.last_dst_samples_loss = [] if self.gan_power != 0: - self.D_src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) + self.GAN_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em) return ( ('src_loss', np.mean(src_loss) ), ('dst_loss', np.mean(dst_loss) ), ) From 8b71b83c940ef63414a9d88780521acba1760974 Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 17 Jul 2021 22:03:59 +0400 Subject: [PATCH 23/30] fix --- mainscripts/Merger.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mainscripts/Merger.py b/mainscripts/Merger.py index fba37f1..0703dc1 100644 --- a/mainscripts/Merger.py +++ b/mainscripts/Merger.py @@ -49,6 +49,7 @@ def main (model_class_name=None, model = models.import_model(model_class_name)(is_training=False, saved_models_path=saved_models_path, force_gpu_idxs=force_gpu_idxs, + force_model_name=force_model_name, cpu_only=cpu_only) predictor_func, predictor_input_shape, cfg = model.get_MergerConfig() From 336e0ca944ba170c43255fbf116cd8ae57114f8c Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 17 Jul 2021 22:04:20 +0400 Subject: [PATCH 24/30] fix --- models/ModelBase.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/models/ModelBase.py b/models/ModelBase.py index 598c19a..f446efa 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -22,6 +22,7 @@ from samplelib import SampleGeneratorBase class ModelBase(object): def __init__(self, is_training=False, + is_exporting=False, saved_models_path=None, training_data_src_path=None, training_data_dst_path=None, @@ -36,6 +37,7 @@ class ModelBase(object): silent_start=False, **kwargs): self.is_training = is_training + self.is_exporting = is_exporting self.saved_models_path = saved_models_path self.training_data_src_path = training_data_src_path self.training_data_dst_path = training_data_dst_path From 9ad9728b4021d1dff62905cce03e2157d0c0868d Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 17 Jul 2021 22:04:34 +0400 Subject: [PATCH 25/30] fix --- mainscripts/ExportDFM.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mainscripts/ExportDFM.py b/mainscripts/ExportDFM.py index 2b0f978..cf7d64e 100644 --- a/mainscripts/ExportDFM.py +++ b/mainscripts/ExportDFM.py @@ -16,7 +16,7 @@ from core.interact import interact as io def main(model_class_name, saved_models_path): model = models.import_model(model_class_name)( - is_training=False, + is_exporting=True, saved_models_path=saved_models_path, cpu_only=True) model.export_dfm () From 493f23c6180d10d7d30f609bb246411ac699391d Mon Sep 17 00:00:00 2001 From: iperov Date: Sat, 17 Jul 2021 23:10:11 +0400 Subject: [PATCH 26/30] update windows magnet link --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index d6a219f..abe18a6 100644 --- a/README.md +++ b/README.md @@ -29,8 +29,8 @@ More than 95% of deepfake videos are created with DeepFaceLab. DeepFaceLab is used by such popular youtube channels as -|![](doc/tiktok_icon.png) [deeptomcruise](https://www.tiktok.com/@deeptomcruise)|![](doc/tiktok_icon.png) [1facerussia](https://www.tiktok.com/@1facerussia)|![](doc/tiktok_icon.png) [arnoldschwarzneggar](https://www.tiktok.com/@arnoldschwarzneggar) -|---|---|---| +|![](doc/tiktok_icon.png) [deeptomcruise](https://www.tiktok.com/@deeptomcruise)|![](doc/tiktok_icon.png) [1facerussia](https://www.tiktok.com/@1facerussia)|![](doc/tiktok_icon.png) [arnoldschwarzneggar](https://www.tiktok.com/@arnoldschwarzneggar)|![](doc/tiktok_icon.png) [mariahcareyathome?](https://www.tiktok.com/@mariahcareyathome?) +|---|---|---|---| |![](doc/youtube_icon.png) [Ctrl Shift Face](https://www.youtube.com/channel/UCKpH0CKltc73e4wh0_pgL3g)|![](doc/youtube_icon.png) [VFXChris Ume](https://www.youtube.com/channel/UCGf4OlX_aTt8DlrgiH3jN3g/videos)|![](doc/youtube_icon.png) [Sham00k](https://www.youtube.com/channel/UCZXbWcv7fSZFTAZV4beckyw/videos)| |---|---|---| @@ -194,7 +194,7 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul -Windows (magnet link) +Windows (magnet link) Last release. Use torrent client to download. From 850b1c5cc994eced9a97d2917323dd636ef6f509 Mon Sep 17 00:00:00 2001 From: iperov Date: Mon, 19 Jul 2021 19:54:24 +0400 Subject: [PATCH 27/30] upd readme: added related works section --- README.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/README.md b/README.md index abe18a6..dcdadc8 100644 --- a/README.md +++ b/README.md @@ -307,6 +307,18 @@ QQ群1095077489 + + + + +
+ +## Related works + +
+neuralchen/SimSwap +swap face using one photo 一张图免训练换脸
+ -
From c5044580c73e21e92133b9b84949ea4b490eeeb0 Mon Sep 17 00:00:00 2001 From: iperov Date: Mon, 19 Jul 2021 19:55:59 +0400 Subject: [PATCH 28/30] upd readme --- README.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index dcdadc8..a8a0a2e 100644 --- a/README.md +++ b/README.md @@ -305,13 +305,12 @@ QQ群1095077489 deepfaker.xyz 中文学习站（非官方)
- - + From 5c7e2f310a06b883e61a831a0bd499e91280a185 Mon Sep 17 00:00:00 2001 From: iperov Date: Mon, 19 Jul 2021 19:58:03 +0400 Subject: [PATCH 29/30] upd readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index a8a0a2e..df7ef67 100644 --- a/README.md +++ b/README.md @@ -313,7 +313,7 @@ QQ群1095077489 +
## Related works +
neuralchen/SimSwap swap face using one photo 一张图免训练换脸
neuralchen/SimSwap -swap face using one photo 一张图免训练换脸
Swapping face using ONE single photo 一张图免训练换脸
From bfa88c5fd9a1615748b15de4999a0f5830d440bd Mon Sep 17 00:00:00 2001 From: iperov Date: Mon, 19 Jul 2021 23:33:56 +0400 Subject: [PATCH 30/30] upd requirements-colab.txt --- requirements-colab.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements-colab.txt b/requirements-colab.txt index 49e8a9f..fccd51e 100644 --- a/requirements-colab.txt +++ b/requirements-colab.txt @@ -1,10 +1,10 @@ tqdm numpy==1.19.3 -h5py==2.9.0 +h5py==2.10.0 opencv-python==4.1.0.25 ffmpeg-python==0.1.17 scikit-image==0.14.2 scipy==1.4.1 colorama -tensorflow-gpu==2.3.1 +tensorflow-gpu==2.4.0 tf2onnx==1.8.4 \ No newline at end of file