mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-19 21:13:20 -07:00
Merge pull request #1 from MachineEditor/JanFschr-bug-fixes
Added tensorboard support, hsv power, bug fixes - removed blur_out_mask
This commit is contained in:
Ognjen
GitHub
commit
0613321ea3
8 changed files with 171 additions and 165 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,7 @@ 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)
|
||||
|
|
@ -147,7 +150,11 @@ 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)]) }
|
||||
|
|
@ -172,13 +179,10 @@ class UNetPatchDiscriminator(nn.ModelBase):
|
|||
|
||||
|
||||
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 = []
|
||||
for conv in self.convs:
|
||||
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)
|
||||
|
|
@ -186,23 +190,20 @@ class UNetPatchDiscriminator(nn.ModelBase):
|
|||
|
||||
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)):
|
||||
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)
|
||||
|
||||
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
|
||||
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):
|
||||
"""
|
||||
|
|
@ -218,7 +219,7 @@ class UNetPatchDiscriminatorV2(nn.ModelBase):
|
|||
ts *= s
|
||||
return rf
|
||||
|
||||
def find_archi(self, target_patch_size, max_layers=6):
|
||||
def find_archi(self, target_patch_size, max_layers=9):
|
||||
"""
|
||||
Find the best configuration of layers using only 3x3 convs for target patch size
|
||||
"""
|
||||
|
|
@ -230,12 +231,12 @@ class UNetPatchDiscriminatorV2(nn.ModelBase):
|
|||
|
||||
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
|
||||
layers.append ( [3, 2])
|
||||
sum_st += 2
|
||||
|
||||
rf = self.calc_receptive_field_size(layers)
|
||||
|
||||
|
|
@ -254,7 +255,7 @@ class UNetPatchDiscriminatorV2(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, use_fp16 = False):
|
||||
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
|
||||
|
||||
|
|
@ -272,11 +273,8 @@ class UNetPatchDiscriminatorV2(nn.ModelBase):
|
|||
|
||||
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)]) }
|
||||
|
||||
|
|
@ -285,12 +283,8 @@ class UNetPatchDiscriminatorV2(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.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)
|
||||
|
|
@ -301,20 +295,18 @@ class UNetPatchDiscriminatorV2(nn.ModelBase):
|
|||
if self.use_fp16:
|
||||
x = tf.cast(x, tf.float16)
|
||||
|
||||
x = tf.nn.leaky_relu( self.in_conv(x), 0.1 )
|
||||
x = tf.nn.leaky_relu( self.in_conv(x), 0.2 )
|
||||
|
||||
encs = []
|
||||
for conv, res in zip(self.convs, self.res):
|
||||
for conv in self.convs:
|
||||
encs.insert(0, x)
|
||||
x = tf.nn.leaky_relu( conv(x), 0.1 )
|
||||
x = res(x)
|
||||
x = tf.nn.leaky_relu( conv(x), 0.2 )
|
||||
|
||||
center_out, x = self.center_out(x), self.center_conv(x)
|
||||
center_out, x = self.center_out(x), tf.nn.leaky_relu( self.center_conv(x), 0.2 )
|
||||
|
||||
for i, (upconv, enc, upres) in enumerate(zip(self.upconvs, encs, self.upres)):
|
||||
x = tf.nn.leaky_relu( upconv(x), 0.1 )
|
||||
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 = upres(x)
|
||||
|
||||
x = self.out_conv(x)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -125,6 +129,25 @@ def trainerThread (s2c, c2s, e,
|
|||
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:
|
||||
previews = model.get_previews()
|
||||
|
|
@ -322,7 +345,6 @@ 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):
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -394,6 +394,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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.')
|
||||
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
|
|
|
|||
|
|
@ -9,3 +9,5 @@ scipy==1.4.1
|
|||
colorama
|
||||
tensorflow-gpu==2.4.0
|
||||
tf2onnx==1.9.3
|
||||
tensorboardX
|
||||
crc32c
|
||||
|
|
|
|||
|
|
@ -12,3 +12,5 @@ pyqt5
|
|||
tf2onnx==1.9.3
|
||||
Flask==1.1.1
|
||||
flask-socketio==4.2.1
|
||||
tensorboardX
|
||||
crc32c
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue