diff --git a/doc/manual_en_google_translated.docx b/doc/manual_en_google_translated.docx index 88a2b90..ac64a7d 100644 Binary files a/doc/manual_en_google_translated.docx and b/doc/manual_en_google_translated.docx differ diff --git a/doc/manual_en_google_translated.pdf b/doc/manual_en_google_translated.pdf index 56e62fa..d6fade9 100644 Binary files a/doc/manual_en_google_translated.pdf and b/doc/manual_en_google_translated.pdf differ diff --git a/doc/manual_ru.pdf b/doc/manual_ru.pdf index bc302e4..2c27e9f 100644 Binary files a/doc/manual_ru.pdf and b/doc/manual_ru.pdf differ diff --git a/doc/manual_ru_source.docx b/doc/manual_ru_source.docx index 1104e45..f37a425 100644 Binary files a/doc/manual_ru_source.docx and b/doc/manual_ru_source.docx differ diff --git a/models/ModelBase.py b/models/ModelBase.py index 9b5b557..ef27065 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -508,6 +508,10 @@ class ModelBase(object): def generate_next_sample(self): return [ generator.generate_next() for generator in self.generator_list] + #overridable + def on_success_train_one_iter(self): + pass + def train_one_iter(self): sample = self.generate_next_sample() iter_time = time.time() @@ -534,7 +538,8 @@ class ModelBase(object): img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8) cv2_imwrite (filepath, img ) - + self.on_success_train_one_iter() + self.iter += 1 return self.iter, iter_time diff --git a/models/Model_SAEv2/Model.py b/models/Model_SAEHD/Model.py similarity index 77% rename from models/Model_SAEv2/Model.py rename to models/Model_SAEHD/Model.py index ceed8f7..bd28693 100644 --- a/models/Model_SAEv2/Model.py +++ b/models/Model_SAEHD/Model.py @@ -30,7 +30,7 @@ class SAEv2Model(ModelBase): self.options['resolution'] = resolution self.options['face_type'] = io.input_str ("Half, mid full, or full face? (h/mf/f, ?:help skip:f) : ", default_face_type, ['h','mf','f'], help_message="Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face.").lower() - self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.") + self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted.") else: self.options['resolution'] = self.options.get('resolution', default_resolution) self.options['face_type'] = self.options.get('face_type', default_face_type) @@ -48,28 +48,18 @@ class SAEv2Model(ModelBase): self.options['archi'] = self.options.get('archi', default_archi) default_ae_dims = 256 if 'liae' in self.options['archi'] else 512 - default_e_ch_dims = 21 - default_d_ch_dims = default_e_ch_dims - def_ca_weights = False + default_ed_ch_dims = 21 if is_first_run: self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, 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['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 ) - default_d_ch_dims = self.options['e_ch_dims'] - self.options['d_ch_dims'] = np.clip ( io.input_int("Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_d_ch_dims) , default_d_ch_dims, help_message="More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU." ), 10, 85 ) - self.options['ca_weights'] = io.input_bool (f"Use CA weights? (y/n, ?:help skip:{yn_str[def_ca_weights]} ) : ", def_ca_weights, help_message="Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run.") + self.options['ed_ch_dims'] = np.clip ( io.input_int("Encoder/Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_ed_ch_dims) , default_ed_ch_dims, 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." ), 10, 85 ) else: self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims) - self.options['e_ch_dims'] = self.options.get('e_ch_dims', default_e_ch_dims) - self.options['d_ch_dims'] = self.options.get('d_ch_dims', default_d_ch_dims) - self.options['ca_weights'] = self.options.get('ca_weights', def_ca_weights) + self.options['ed_ch_dims'] = self.options.get('ed_ch_dims', default_ed_ch_dims) default_face_style_power = 0.0 default_bg_style_power = 0.0 if is_first_run or ask_override: - def_pixel_loss = self.options.get('pixel_loss', False) - self.options['pixel_loss'] = io.input_bool (f"Use pixel loss? (y/n, ?:help skip:{yn_str[def_pixel_loss]} ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. Enabling this option too early increases the chance of model collapse.") - default_face_style_power = default_face_style_power if is_first_run else self.options.get('face_style_power', default_face_style_power) self.options['face_style_power'] = np.clip ( io.input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_face_style_power), default_face_style_power, help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 ) @@ -82,8 +72,7 @@ class SAEv2Model(ModelBase): self.options['apply_random_ct'] = io.input_bool (f"Apply random color transfer to src faceset? (y/n, ?:help skip:{yn_str[default_apply_random_ct]}) : ", default_apply_random_ct, help_message="Increase variativity of src samples by apply LCT color transfer from random dst samples. It is like 'face_style' learning, but more precise color transfer and without risk of model collapse, also it does not require additional GPU resources, but the training time may be longer, due to the src faceset is becoming more diverse.") default_true_face_training = False if is_first_run else self.options.get('true_face_training', False) - self.options['true_face_training'] = io.input_bool (f"Enable 'true face' training? (y/n, ?:help skip:{yn_str[default_true_face_training]}) : ", default_true_face_training, help_message="Result face will be more like src. Enable it only after 100k iters, when face is sharp enough.") - + self.options['true_face_training'] = io.input_bool (f"Enable 'true face' training? (y/n, ?:help skip:{yn_str[default_true_face_training]}) : ", default_true_face_training, help_message="The result face will be more like src and will get extra sharpness. Enable it for last 10-20k iterations before conversion.") if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301 default_clipgrad = False if is_first_run else self.options.get('clipgrad', False) @@ -92,7 +81,6 @@ class SAEv2Model(ModelBase): self.options['clipgrad'] = False else: - self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power) self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power) self.options['apply_random_ct'] = self.options.get('apply_random_ct', False) @@ -112,8 +100,7 @@ class SAEv2Model(ModelBase): learn_mask = self.options['learn_mask'] ae_dims = self.options['ae_dims'] - e_ch_dims = self.options['e_ch_dims'] - d_ch_dims = self.options['d_ch_dims'] + ed_ch_dims = self.options['ed_ch_dims'] self.pretrain = self.options['pretrain'] = self.options.get('pretrain', False) if not self.pretrain: self.options.pop('pretrain') @@ -136,9 +123,12 @@ class SAEv2Model(ModelBase): lowest_dense_res = resolution // 16 e_dims = output_nc*e_ch_dims - def downscale (dim, kernel_size=5, dilation_rate=1): + def downscale (dim, kernel_size=5, dilation_rate=1, use_activator=True): def func(x): - return SubpixelDownscaler()(LeakyReLU(0.1)(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x))) + if not use_activator: + return SubpixelDownscaler()(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x)) + else: + return SubpixelDownscaler()(LeakyReLU(0.1)(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x))) return func def upscale (dim, size=(2,2)): @@ -146,27 +136,31 @@ class SAEv2Model(ModelBase): return SubpixelUpscaler(size=size)(LeakyReLU(0.1)(Conv2D(dim * np.prod(size) , kernel_size=3, strides=1, padding='same')(x))) return func - def enc_flow(e_dims, ae_dims, lowest_dense_res): + def enc_flow(e_ch_dims, ae_dims, lowest_dense_res): + dims = output_nc * e_ch_dims + if dims % 2 != 0: + dims += 1 + def func(inp): - x = downscale(e_dims , 3, 1 )(inp) - x = downscale(e_dims*2, 3, 1 )(x) - x = downscale(e_dims*4, 3, 1 )(x) - x0 = downscale(e_dims*8, 3, 1 )(x) + x = downscale(dims , 3, 1 )(inp) + x = downscale(dims*2, 3, 1 )(x) + x = downscale(dims*4, 3, 1 )(x) + x0 = downscale(dims*8, 3, 1 )(x) - x = downscale(e_dims , 5, 1 )(inp) - x = downscale(e_dims*2, 5, 1 )(x) - x = downscale(e_dims*4, 5, 1 )(x) - x1 = downscale(e_dims*8, 5, 1 )(x) + x = downscale(dims , 5, 1 )(inp) + x = downscale(dims*2, 5, 1 )(x) + x = downscale(dims*4, 5, 1 )(x) + x1 = downscale(dims*8, 5, 1 )(x) - x = downscale(e_dims , 5, 2 )(inp) - x = downscale(e_dims*2, 5, 2 )(x) - x = downscale(e_dims*4, 5, 2 )(x) - x2 = downscale(e_dims*8, 5, 2 )(x) + x = downscale(dims , 5, 2 )(inp) + x = downscale(dims*2, 5, 2 )(x) + x = downscale(dims*4, 5, 2 )(x) + x2 = downscale(dims*8, 5, 2 )(x) - x = downscale(e_dims , 7, 2 )(inp) - x = downscale(e_dims*2, 7, 2 )(x) - x = downscale(e_dims*4, 7, 2 )(x) - x3 = downscale(e_dims*8, 7, 2 )(x) + x = downscale(dims , 7, 2 )(inp) + x = downscale(dims*2, 7, 2 )(x) + x = downscale(dims*4, 7, 2 )(x) + x3 = downscale(dims*8, 7, 2 )(x) x = Concatenate()([x0,x1,x2,x3]) @@ -177,8 +171,11 @@ class SAEv2Model(ModelBase): return x return func - def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True): + def dec_flow(output_nc, d_ch_dims, is_mask=False): dims = output_nc * d_ch_dims + if dims % 2 != 0: + dims += 1 + def ResidualBlock(dim): def func(inp): x = Conv2D(dim, kernel_size=3, padding='same')(inp) @@ -188,64 +185,33 @@ class SAEv2Model(ModelBase): x = LeakyReLU(0.2)(x) return x return func - """ - def func(x): - x0 = upscale(dims, size=(8,8) )(x) - - x = upscale(dims*4)(x) - - if add_residual_blocks: - x = ResidualBlock(dims*4)(x) - - x1 = upscale(dims, size=(4,4) )(x) - - x = upscale(dims*2)(x) - - if add_residual_blocks: - x = ResidualBlock(dims*2)(x) - - x2 = upscale(dims, size=(2,2) )(x) - - x = upscale(dims)(x) - - if add_residual_blocks: - x = ResidualBlock(dims)(x) - - x = Add()([x0,x1,x2,x]) - - - return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x) - """ + def func(x): - x = upscale(dims*8)(x) - - if add_residual_blocks: - x = ResidualBlock(dims*8)(x) - - x = upscale(dims*4)(x) - - if add_residual_blocks: - x = ResidualBlock(dims*4)(x) - - x = upscale(dims*2)(x) - - if add_residual_blocks: - x = ResidualBlock(dims*2)(x) - + for i in [8,4,2]: + x = upscale(dims*i)(x) + + if not is_mask: + x0 = x + x = upscale( (dims*i)//2 )(x) + x = ResidualBlock( (dims*i)//2 )(x) + x = downscale( dims*i, use_activator=False ) (x) + x = Add()([x, x0]) + x = LeakyReLU(0.2)(x) + return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x) - + return func - self.encoder = modelify(enc_flow(e_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) ) + self.encoder = modelify(enc_flow(e_ch_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) ) sh = K.int_shape( self.encoder.outputs[0] )[1:] self.decoder_src = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) ) self.decoder_dst = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) ) if learn_mask: - self.decoder_srcm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) ) - self.decoder_dstm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) ) + self.decoder_srcm = modelify(dec_flow(1, d_ch_dims, is_mask=True)) ( Input(sh) ) + self.decoder_dstm = modelify(dec_flow(1, d_ch_dims, is_mask=True)) ( Input(sh) ) self.src_dst_trainable_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights @@ -285,14 +251,15 @@ class SAEv2Model(ModelBase): output_nc = 3 bgr_shape = (resolution, resolution, output_nc) mask_shape = (resolution, resolution, 1) - - e_dims = output_nc*e_ch_dims - + lowest_dense_res = resolution // 16 - def downscale (dim, kernel_size=5, dilation_rate=1): + def downscale (dim, kernel_size=5, dilation_rate=1, use_activator=True): def func(x): - return SubpixelDownscaler()(LeakyReLU(0.1)(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x))) + if not use_activator: + return SubpixelDownscaler()(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x)) + else: + return SubpixelDownscaler()(LeakyReLU(0.1)(Conv2D(dim // 4, kernel_size=kernel_size, strides=1, dilation_rate=dilation_rate, padding='same')(x))) return func def upscale (dim): @@ -300,27 +267,31 @@ class SAEv2Model(ModelBase): return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x))) return func - def enc_flow(e_dims): - def func(inp): - x = downscale(e_dims , 3, 1 )(inp) - x = downscale(e_dims*2, 3, 1 )(x) - x = downscale(e_dims*4, 3, 1 )(x) - x0 = downscale(e_dims*8, 3, 1 )(x) + def enc_flow(e_ch_dims): + dims = output_nc*e_ch_dims + if dims % 2 != 0: + dims += 1 - x = downscale(e_dims , 5, 1 )(inp) - x = downscale(e_dims*2, 5, 1 )(x) - x = downscale(e_dims*4, 5, 1 )(x) - x1 = downscale(e_dims*8, 5, 1 )(x) + def func(inp): + x = downscale(dims , 3, 1 )(inp) + x = downscale(dims*2, 3, 1 )(x) + x = downscale(dims*4, 3, 1 )(x) + x0 = downscale(dims*8, 3, 1 )(x) - x = downscale(e_dims , 5, 2 )(inp) - x = downscale(e_dims*2, 5, 2 )(x) - x = downscale(e_dims*4, 5, 2 )(x) - x2 = downscale(e_dims*8, 5, 2 )(x) + x = downscale(dims , 5, 1 )(inp) + x = downscale(dims*2, 5, 1 )(x) + x = downscale(dims*4, 5, 1 )(x) + x1 = downscale(dims*8, 5, 1 )(x) - x = downscale(e_dims , 7, 2 )(inp) - x = downscale(e_dims*2, 7, 2 )(x) - x = downscale(e_dims*4, 7, 2 )(x) - x3 = downscale(e_dims*8, 7, 2 )(x) + x = downscale(dims , 5, 2 )(inp) + x = downscale(dims*2, 5, 2 )(x) + x = downscale(dims*4, 5, 2 )(x) + x2 = downscale(dims*8, 5, 2 )(x) + + x = downscale(dims , 7, 2 )(inp) + x = downscale(dims*2, 7, 2 )(x) + x = downscale(dims*4, 7, 2 )(x) + x3 = downscale(dims*8, 7, 2 )(x) x = Concatenate()([x0,x1,x2,x3]) @@ -337,38 +308,39 @@ class SAEv2Model(ModelBase): return x return func - def dec_flow(output_nc, d_ch_dims, add_residual_blocks=True): - d_dims = output_nc*d_ch_dims + def dec_flow(output_nc, d_ch_dims, is_mask=False): + dims = output_nc * d_ch_dims + if dims % 2 != 0: + dims += 1 + def ResidualBlock(dim): def func(inp): x = Conv2D(dim, kernel_size=3, padding='same')(inp) x = LeakyReLU(0.2)(x) - x = Conv2D(dim, kernel_size=3, padding='same')(inp) + x = Conv2D(dim, kernel_size=3, padding='same')(x) x = Add()([x, inp]) x = LeakyReLU(0.2)(x) return x return func - + def func(x): - x = upscale(d_dims*8)(x) - - if add_residual_blocks: - x = ResidualBlock(d_dims*8)(x) + + for i in [8,4,2]: + x = upscale(dims*i)(x) - x = upscale(d_dims*4)(x) - - if add_residual_blocks: - x = ResidualBlock(d_dims*4)(x) - - x = upscale(d_dims*2)(x) - - if add_residual_blocks: - x = ResidualBlock(d_dims*2)(x) - + if not is_mask: + x0 = x + x = upscale( (dims*i)//2 )(x) + x = ResidualBlock( (dims*i)//2 )(x) + x = downscale( dims*i, use_activator=False ) (x) + x = Add()([x, x0]) + x = LeakyReLU(0.2)(x) + return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x) + return func - self.encoder = modelify(enc_flow(e_dims)) ( Input(bgr_shape) ) + self.encoder = modelify(enc_flow(e_ch_dims)) ( Input(bgr_shape) ) sh = K.int_shape( self.encoder.outputs[0] )[1:] self.inter_B = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) ) @@ -378,7 +350,7 @@ class SAEv2Model(ModelBase): self.decoder = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) ) if learn_mask: - self.decoderm = modelify(dec_flow(1, d_ch_dims, add_residual_blocks=False)) ( Input(sh) ) + self.decoderm = modelify(dec_flow(1, d_ch_dims, is_mask=True)) ( Input(sh) ) self.src_dst_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights @@ -424,9 +396,9 @@ class SAEv2Model(ModelBase): return ar if 'df' in self.options['archi']: - self.model = SAEDFModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask) + self.model = SAEDFModel (resolution, ae_dims, ed_ch_dims, ed_ch_dims, learn_mask) elif 'liae' in self.options['archi']: - self.model = SAELIAEModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask) + self.model = SAELIAEModel (resolution, ae_dims, ed_ch_dims, ed_ch_dims, learn_mask) self.opt_dis_model = [] @@ -467,18 +439,13 @@ class SAEv2Model(ModelBase): if not self.is_first_run(): loaded, not_loaded = self.load_weights_safe(not_loaded) - CA_models = [] - if self.options.get('ca_weights', False): - CA_models += [ model for model, _ in not_loaded ] + CA_models = [ model for model, _ in not_loaded ] - CA_conv_weights_list = [] + self.CA_conv_weights_list = [] for model in CA_models: for layer in model.layers: if type(layer) == keras.layers.Conv2D: - CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights - - if len(CA_conv_weights_list) != 0: - CAInitializerMP ( CA_conv_weights_list ) + self.CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights target_srcm = gaussian_blur( max(1, resolution // 32) )(self.model.target_srcm) target_dstm = gaussian_blur( max(1, resolution // 32) )(self.model.target_dstm) @@ -501,10 +468,8 @@ class SAEv2Model(ModelBase): self.src_dst_mask_opt = RMSprop(lr=5e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) self.D_opt = RMSprop(lr=5e-5, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1) - if not self.options['pixel_loss']: - src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) ) - else: - src_loss = K.mean ( 50*K.square( target_src_masked_opt - pred_src_src_masked_opt ) ) + src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) ) + src_loss += K.mean ( 10*K.square( target_src_masked_opt - pred_src_src_masked_opt ) ) face_style_power = self.options['face_style_power'] / 100.0 if face_style_power != 0: @@ -512,15 +477,11 @@ class SAEv2Model(ModelBase): bg_style_power = self.options['bg_style_power'] / 100.0 if bg_style_power != 0: - if not self.options['pixel_loss']: - src_loss += K.mean( (10*bg_style_power)*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( psd_target_dst_anti_masked, target_dst_anti_masked )) - else: - src_loss += K.mean( (50*bg_style_power)*K.square( psd_target_dst_anti_masked - target_dst_anti_masked )) + src_loss += K.mean( (10*bg_style_power)*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( psd_target_dst_anti_masked, target_dst_anti_masked )) + src_loss += K.mean( (10*bg_style_power)*K.square( psd_target_dst_anti_masked - target_dst_anti_masked )) - if not self.options['pixel_loss']: - dst_loss = K.mean( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)(target_dst_masked_opt, pred_dst_dst_masked_opt) ) - else: - dst_loss = K.mean( 50*K.square( target_dst_masked_opt - pred_dst_dst_masked_opt ) ) + dst_loss = K.mean( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)(target_dst_masked_opt, pred_dst_dst_masked_opt) ) + dst_loss += K.mean( 10*K.square( target_dst_masked_opt - pred_dst_dst_masked_opt ) ) opt_D_loss = [] if self.true_face_training: @@ -532,7 +493,7 @@ class SAEv2Model(ModelBase): src_code_d_zeros = K.zeros_like(src_code_d) dst_code_d = self.dis( self.model.dst_code ) dst_code_d_ones = K.ones_like(dst_code_d) - opt_D_loss = [ DLoss(src_code_d_ones, src_code_d) ] + opt_D_loss = [ 0.01*DLoss(src_code_d_ones, src_code_d) ] loss_D = (DLoss(dst_code_d_ones , dst_code_d) + \ DLoss(src_code_d_zeros, src_code_d) ) * 0.5 @@ -606,6 +567,13 @@ class SAEv2Model(ModelBase): #override def onSave(self): self.save_weights_safe( self.get_model_filename_list()+self.opt_dis_model ) + + #override + def on_success_train_one_iter(self): + if len(self.CA_conv_weights_list) != 0: + exec(nnlib.import_all(), locals(), globals()) + CAInitializerMP ( self.CA_conv_weights_list ) + self.CA_conv_weights_list = [] #override def onTrainOneIter(self, generators_samples, generators_list): @@ -645,7 +613,7 @@ class SAEv2Model(ModelBase): st.append ( np.concatenate ( ar, axis=1) ) - result += [ ('SAE', np.concatenate (st, axis=0 )), ] + result += [ ('SAEHD', np.concatenate (st, axis=0 )), ] if self.options['learn_mask']: st_m = [] @@ -653,7 +621,7 @@ class SAEv2Model(ModelBase): ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) st_m.append ( np.concatenate ( ar, axis=1) ) - result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ] + result += [ ('SAEHD masked', np.concatenate (st_m, axis=0 )), ] return result diff --git a/models/Model_SAEv2/__init__.py b/models/Model_SAEHD/__init__.py similarity index 100% rename from models/Model_SAEv2/__init__.py rename to models/Model_SAEHD/__init__.py