mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-05 20:42:11 -07:00
removed SAEv2,
added SAEHD model ( High Definition Styled AutoEncoder ) This is a new heavyweight model for high-end cards to achieve maximum possible deepfake quality in 2020. Differences from SAE: + new encoder produces more stable face and less scale jitter before: https://i.imgur.com/4jUcol8.gifv after: https://i.imgur.com/lyiax49.gifv - scale of the face is less changed within frame size + new decoder produces subpixel clear result + pixel loss and dssim loss are merged together to achieve both training speed and pixel trueness + by default networks will be initialized with CA weights, but only after first successful iteration therefore you can test network size and batch size before weights initialization process + new neural network optimizer consumes less VRAM than before + added option <Enable 'true face' training> The result face will be more like src and will get extra sharpness. example: https://i.imgur.com/ME3A7dI.gifv Enable it for last 15-30k iterations before conversion. + encoder and decoder dims are merged to one parameter encoder/decoder dims + added mid-full face, which covers 30% more area than half face.
This commit is contained in:
Colombo
parent
cbc18b2d41
commit
ac7725163d
7 changed files with 128 additions and 155 deletions
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
Loading…
Add table
Add a link
Reference in a new issue