diff --git a/core/imagelib/warp.py b/core/imagelib/warp.py index 4f6f68a..4de56d2 100644 --- a/core/imagelib/warp.py +++ b/core/imagelib/warp.py @@ -104,7 +104,7 @@ def gen_pts(W, H, rnd_state=None): return pts1, pts2 -def gen_warp_params (w, flip=False, rotation_range=[-2,2], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], rnd_state=None ): +def gen_warp_params (w, flip=False, rotation_range=[-2,2], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], rnd_state=None, warp_rnd_state=None): if rnd_state is None: rnd_state = np.random if warp_rnd_state is None: @@ -178,4 +178,4 @@ def warp_by_params (params, img, can_warp, can_transform, can_flip, border_repli img = img[...,None] if can_flip and params['flip']: img = img[:,::-1,...] - return img \ No newline at end of file + return img diff --git a/core/leras/models/PatchDiscriminator.py b/core/leras/models/PatchDiscriminator.py index 63dd2b5..66ab02e 100644 --- a/core/leras/models/PatchDiscriminator.py +++ b/core/leras/models/PatchDiscriminator.py @@ -76,9 +76,11 @@ class PatchDiscriminator(nn.ModelBase): nn.PatchDiscriminator = PatchDiscriminator + class UNetPatchDiscriminator(nn.ModelBase): """ Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks" + Based on iperov commit 11add4cd4f5a61df26a8659f4cc5c8d9467bf5f8 from Jan 3, 2021 with added fp16 option """ def calc_receptive_field_size(self, layers): """ @@ -133,6 +135,130 @@ class UNetPatchDiscriminator(nn.ModelBase): 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', 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) + x = tf.nn.leaky_relu(x, 0.2) + x = self.conv2(x) + x = tf.nn.leaky_relu(inp + x, 0.2) + return x + + 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)]) } + + 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', dtype=conv_dtype) ) + + 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', dtype=conv_dtype) ) + + 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', dtype=conv_dtype) + + 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): + x = tf.nn.leaky_relu( self.in_conv(x), 0.2 ) + + encs = [] + for conv, res1,res2 in zip(self.convs, self.res1, self.res2): + 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)): + 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) + +nn.UNetPatchDiscriminator = UNetPatchDiscriminator + +class UNetPatchDiscriminatorV2(nn.ModelBase): + """ + Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks" + Based on iperov commit 35877dbfd724c22040f421e93c1adbb7142e5b5d from Jul 14, 2021 + """ + def calc_receptive_field_size(self, layers): + """ + result the same as https://fomoro.com/research/article/receptive-field-calculatorindex.html + """ + rf = 0 + ts = 1 + for i, (k, s) in enumerate(layers): + if i == 0: + rf = k + else: + rf += (k-1)*ts + ts *= s + return rf + + def find_archi(self, target_patch_size, max_layers=9): + """ + Find the best configuration of layers using only 3x3 convs for target patch size + """ + s = {} + for layers_count in range(1,max_layers+1): + val = 1 << (layers_count-1) + while True: + val -= 1 + + layers = [] + sum_st = 0 + layers.append ( [3, 2]) + sum_st += 2 + for i in range(layers_count-1): + st = 1 + (1 if val & (1 << i) !=0 else 0 ) + layers.append ( [3, st ]) + sum_st += st + + rf = self.calc_receptive_field_size(layers) + + s_rf = s.get(rf, None) + if s_rf is None: + s[rf] = (layers_count, sum_st, layers) + else: + if layers_count < s_rf[0] or \ + ( layers_count == s_rf[0] and sum_st > s_rf[1] ): + s[rf] = (layers_count, sum_st, layers) + + if val == 0: + break + + x = sorted(list(s.keys())) + 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, 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', dtype=conv_dtype) @@ -149,7 +275,7 @@ class UNetPatchDiscriminator(nn.ModelBase): self.convs = [] self.upconvs = [] 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)]) } self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID', dtype=conv_dtype) @@ -157,14 +283,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', dtype=conv_dtype) ) - 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', dtype=conv_dtype) ) - 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', dtype=conv_dtype) self.center_out = nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID', dtype=conv_dtype) @@ -181,9 +301,7 @@ class UNetPatchDiscriminator(nn.ModelBase): 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) in enumerate(zip(self.upconvs, encs)): @@ -197,131 +315,5 @@ class UNetPatchDiscriminator(nn.ModelBase): x = tf.cast(x, tf.float32) return center_out, x - -nn.UNetPatchDiscriminator = UNetPatchDiscriminator - -class UNetPatchDiscriminatorV2(nn.ModelBase): - """ - Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks" - """ - def calc_receptive_field_size(self, layers): - """ - result the same as https://fomoro.com/research/article/receptive-field-calculatorindex.html - """ - rf = 0 - ts = 1 - for i, (k, s) in enumerate(layers): - if i == 0: - rf = k - else: - rf += (k-1)*ts - ts *= s - return rf - - def find_archi(self, target_patch_size, max_layers=6): - """ - Find the best configuration of layers using only 3x3 convs for target patch size - """ - s = {} - for layers_count in range(1,max_layers+1): - val = 1 << (layers_count-1) - while True: - val -= 1 - - layers = [] - sum_st = 0 - for i in range(layers_count-1): - st = 1 + (1 if val & (1 << i) !=0 else 0 ) - layers.append ( [3, st ]) - sum_st += st - layers.append ( [3, 2]) - sum_st += 2 - - rf = self.calc_receptive_field_size(layers) - - s_rf = s.get(rf, None) - if s_rf is None: - s[rf] = (layers_count, sum_st, layers) - else: - if layers_count < s_rf[0] or \ - ( layers_count == s_rf[0] and sum_st > s_rf[1] ): - s[rf] = (layers_count, sum_st, layers) - - if val == 0: - break - - x = sorted(list(s.keys())) - q=x[np.abs(np.array(x)-target_patch_size).argmin()] - return s[q][2] - - def on_build(self, patch_size, in_ch, 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', 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) - x = tf.nn.leaky_relu(x, 0.2) - x = self.conv2(x) - x = tf.nn.leaky_relu(inp + x, 0.2) - return x - - prev_ch = in_ch - self.convs = [] - self.res = [] - self.upconvs = [] - self.upres = [] - layers = self.find_archi(patch_size) - base_ch = 16 - - 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', 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', dtype=conv_dtype) ) - - self.res.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', dtype=conv_dtype) ) - - self.upres.insert (0, ResidualBlock(level_chs[i-1]*2) ) - - 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', 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.1 ) - - encs = [] - for conv, res in zip(self.convs, self.res): - encs.insert(0, x) - x = tf.nn.leaky_relu( conv(x), 0.1 ) - x = res(x) - - center_out, x = self.center_out(x), self.center_conv(x) - - for i, (upconv, enc, upres) in enumerate(zip(self.upconvs, encs, self.upres)): - x = tf.nn.leaky_relu( upconv(x), 0.1 ) - x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis) - x = upres(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.UNetPatchDiscriminatorV2 = UNetPatchDiscriminatorV2 diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py index 732d6f9..f8030da 100644 --- a/mainscripts/Trainer.py +++ b/mainscripts/Trainer.py @@ -25,13 +25,17 @@ class TensorBoardTool: def run(self): from tensorboard import default from tensorboard import program + from tensorboard import version as tb_version # remove http messages log = logging.getLogger('werkzeug').setLevel(logging.ERROR) # Start tensorboard server tb = program.TensorBoard(default.get_plugins()) - tb.configure(argv=[None, '--logdir', self.dir_path, '--port', '6006', '--bind_all']) + tb_argv = [None, '--logdir', self.dir_path, '--port', '6006'] + if int(tb_version.VERSION[0])>=2: + tb_argv.append("--bind_all") + tb.configure(argv=tb_argv) url = tb.launch() - print('Launched TensorBoard at {}'.format(url)) + io.log_info('Launched TensorBoard at {}\n'.format(url)) def process_img_for_tensorboard(input_img): # convert format from bgr to rgb @@ -124,6 +128,25 @@ def trainerThread (s2c, c2s, e, def model_backup(): if not debug and not is_reached_goal: model.create_backup() + + def log_step(step, step_time, src_loss, dst_loss): + c2s.put({ + 'op': 'tb', + 'action': 'step', + 'step': step, + 'step_time': step_time, + 'src_loss': src_loss, + 'dst_loss': dst_loss + }) + + def log_previews(step, previews, static_previews): + c2s.put({ + 'op': 'tb', + 'action': 'preview', + 'step': step, + 'previews': previews, + 'static_previews': static_previews + }) def send_preview(): if not debug: @@ -322,8 +345,7 @@ def handle_tensorboard_op(input): log_tensorboard_previews(step, previews, 'preview', train_summary_writer) if static_previews is not None: log_tensorboard_previews(step, static_previews, 'static_preview', train_summary_writer) - c2s.put({'op': 'close'}) - + class Zoom(Enum): ZOOM_25 = (1 / 4, '25%') @@ -584,7 +606,7 @@ def main(**kwargs): selected_preview = selected_preview % len(previews) update_preview = True elif op == 'tb': - handle_tensorboard_op(input) + handle_tensorboard_op(item) elif op == 'close': break diff --git a/models/ModelBase.py b/models/ModelBase.py index 634e797..44d872e 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -393,6 +393,9 @@ 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) diff --git a/models/Model_AMP/Model.py b/models/Model_AMP/Model.py index 23050e9..da49054 100644 --- a/models/Model_AMP/Model.py +++ b/models/Model_AMP/Model.py @@ -51,6 +51,7 @@ class AMPModel(ModelBase): default_ct_mode = self.options['ct_mode'] = self.load_or_def_option('ct_mode', 'none') default_random_color = self.options['random_color'] = self.load_or_def_option('random_color', False) default_clipgrad = self.options['clipgrad'] = self.load_or_def_option('clipgrad', False) + default_use_fp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) ask_override = self.ask_override() if self.is_first_run() or ask_override: @@ -167,7 +168,7 @@ class AMPModel(ModelBase): adabelief = self.options['adabelief'] - # use_fp16 = self.options['use_fp16'] + use_fp16 = self.options['use_fp16'] if self.is_exporting: use_fp16 = io.input_bool ("Export quantized?", False, help_message='Makes the exported model faster. If you have problems, disable this option.') diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index e922705..feb2a49 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -68,6 +68,7 @@ class SAEHDModel(ModelBase): default_random_color = self.options['random_color'] = self.load_or_def_option('random_color', False) 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_use_fp16 = self.options['use_fp16'] = self.load_or_def_option('use_fp16', False) ask_override = self.ask_override() if self.is_first_run() or ask_override: @@ -151,6 +152,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.options['blur_out_mask'] = io.input_bool ("Blur out mask", default_blur_out_mask, help_message='Blurs nearby area outside of applied face mask of training samples. The result is the background near the face is smoothed and less noticeable on swapped face. The exact xseg mask in src and dst faceset is required.') default_gan_power = self.options['gan_power'] = self.load_or_def_option('gan_power', 0.0) + default_gan_version = self.options['gan_version'] = self.load_or_def_option('gan_version', 2) default_gan_patch_size = self.options['gan_patch_size'] = self.load_or_def_option('gan_patch_size', self.options['resolution'] // 8) default_gan_dims = self.options['gan_dims'] = self.load_or_def_option('gan_dims', 16) default_gan_smoothing = self.options['gan_smoothing'] = self.load_or_def_option('gan_smoothing', 0.1) @@ -178,6 +180,8 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 10.0", help_message="Train the network in Generative Adversarial manner. Forces the neural network to learn small details of the face. Enable it only when the face is trained enough and don't disable. Typical value is 0.1"), 0.0, 10.0 ) if self.options['gan_power'] != 0.0: + self.options['gan_version'] = np.clip (io.input_int("GAN version", default_gan_version, add_info="2 or 3", help_message="Choose GAN version (v2: 7/16/2020, v3: 1/3/2021):"), 2, 3) + if self.options['gan_version'] == 3: 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 @@ -253,10 +257,6 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.is_exporting: use_fp16 = io.input_bool ("Export quantized?", False, help_message='Makes the exported model faster. If you have problems, disable this option.') - use_fp16 = False - if self.is_exporting: - use_fp16 = io.input_bool ("Export quantized?", False, help_message='Makes the exported model faster. If you have problems, disable this option.') - 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 @@ -440,22 +440,6 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... y = tf.where(tf.equal(y, 0), tf.ones_like(y), y) gpu_target_dst = gpu_target_dst*gpu_target_dstm_all + (x/y)*gpu_target_dstm_anti - gpu_target_srcm_anti = 1-gpu_target_srcm - gpu_target_dstm_anti = 1-gpu_target_dstm - - if blur_out_mask: - sigma = resolution / 128 - - x = nn.gaussian_blur(gpu_target_src*gpu_target_srcm_anti, sigma) - y = 1-nn.gaussian_blur(gpu_target_srcm, sigma) - y = tf.where(tf.equal(y, 0), tf.ones_like(y), y) - gpu_target_src = gpu_target_src*gpu_target_srcm + (x/y)*gpu_target_srcm_anti - - x = nn.gaussian_blur(gpu_target_dst*gpu_target_dstm_anti, sigma) - y = 1-nn.gaussian_blur(gpu_target_dstm, sigma) - y = tf.where(tf.equal(y, 0), tf.ones_like(y), y) - gpu_target_dst = gpu_target_dst*gpu_target_dstm + (x/y)*gpu_target_dstm_anti - # process model tensors if 'df' in archi_type: @@ -822,7 +806,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... 'random_noise': self.options['random_noise'], 'random_blur': self.options['random_blur'], 'random_jpeg': self.options['random_jpeg'], - 'transform':True, 'channel_type' : channel_type, 'ct_mode': ct_mode, + 'transform':True, 'channel_type' : channel_type, 'ct_mode': ct_mode, 'random_hsv_shift_amount' : random_hsv_power, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : channel_type, '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}, diff --git a/requirements-colab.txt b/requirements-colab.txt index 1d0f166..9d638f5 100644 --- a/requirements-colab.txt +++ b/requirements-colab.txt @@ -9,3 +9,5 @@ scipy==1.4.1 colorama tensorflow-gpu==2.4.0 tf2onnx==1.9.3 +tensorboardX +crc32c diff --git a/requirements-cuda.txt b/requirements-cuda.txt index 41bcc29..1cc4902 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -12,3 +12,5 @@ pyqt5 tf2onnx==1.9.3 Flask==1.1.1 flask-socketio==4.2.1 +tensorboardX +crc32c