github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-08-19 13:09:56 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

remove all trailing spaces pt2

Browse source
This commit is contained in:
camjac251 2021-07-24 10:13:39 -05:00
commit b2f287e52e
7 changed files with 112 additions and 112 deletions
Download patch file Download diff file Expand all files Collapse all files

26
core/leras/archis/DeepFakeArchi.py
View file

@ -16,7 +16,7 @@ class DeepFakeArchi(nn.ArchiBase):
conv_dtype = tf.float16 if use_fp16 else tf.float32 conv_dtype = tf.float16 if use_fp16 else tf.float32
if mod is None: if mod is None:
class Downscale(nn.ModelBase): class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, *kwargs ): def __init__(self, in_ch, out_ch, kernel_size=5, *kwargs ):
@ -79,8 +79,8 @@ class DeepFakeArchi(nn.ArchiBase):
self.in_ch = in_ch self.in_ch = in_ch
self.e_ch = e_ch self.e_ch = e_ch
super().__init__(**kwargs) super().__init__(**kwargs)
def on_build(self): def on_build(self):
self.down1 = DownscaleBlock(self.in_ch, self.e_ch, n_downscales=4, kernel_size=5) self.down1 = DownscaleBlock(self.in_ch, self.e_ch, n_downscales=4, kernel_size=5)
def forward(self, x): def forward(self, x):
@ -90,10 +90,10 @@ class DeepFakeArchi(nn.ArchiBase):
if use_fp16: if use_fp16:
x = tf.cast(x, tf.float32) x = tf.cast(x, tf.float32)
return x return x
def get_out_res(self, res): def get_out_res(self, res):
return res // (2**4) return res // (2**4)
def get_out_ch(self): def get_out_ch(self):
return self.e_ch * 8 return self.e_ch * 8
@ -106,10 +106,10 @@ class DeepFakeArchi(nn.ArchiBase):
def on_build(self): def on_build(self):
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
if 'u' in opts: if 'u' in opts:
self.dense_norm = nn.DenseNorm() self.dense_norm = nn.DenseNorm()
self.dense1 = nn.Dense( in_ch, ae_ch ) self.dense1 = nn.Dense( in_ch, ae_ch )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_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) self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
@ -121,7 +121,7 @@ class DeepFakeArchi(nn.ArchiBase):
x = self.dense1(x) x = self.dense1(x)
x = self.dense2(x) 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, self.ae_out_ch)
if use_fp16: if use_fp16:
x = tf.cast(x, tf.float16) x = tf.cast(x, tf.float16)
x = self.upscale1(x) x = self.upscale1(x)
@ -134,7 +134,7 @@ class DeepFakeArchi(nn.ArchiBase):
return self.ae_out_ch return self.ae_out_ch
class Decoder(nn.ModelBase): 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.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
self.upscale1 = Upscale(d_ch*8, d_ch*4, 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) self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
@ -182,13 +182,13 @@ class DeepFakeArchi(nn.ArchiBase):
if 'd' in opts: if 'd' in opts:
m = self.upscalem3(m) m = self.upscalem3(m)
m = tf.nn.sigmoid(self.out_convm(m)) m = tf.nn.sigmoid(self.out_convm(m))
if use_fp16: if use_fp16:
x = tf.cast(x, tf.float32) x = tf.cast(x, tf.float32)
m = tf.cast(m, tf.float32) m = tf.cast(m, tf.float32)
return x, m return x, m
self.Encoder = Encoder self.Encoder = Encoder
self.Inter = Inter self.Inter = Inter
self.Decoder = Decoder self.Decoder = Decoder

14
core/leras/models/PatchDiscriminator.py
View file

@ -111,7 +111,7 @@ class UNetPatchDiscriminator(nn.ModelBase):
for i in range(layers_count-1): for i in range(layers_count-1):
st = 1 + (1 if val & (1 << i) !=0 else 0 ) st = 1 + (1 if val & (1 << i) !=0 else 0 )
layers.append ( [3, st ]) layers.append ( [3, st ])
sum_st += st sum_st += st
rf = self.calc_receptive_field_size(layers) rf = self.calc_receptive_field_size(layers)
@ -132,8 +132,8 @@ class UNetPatchDiscriminator(nn.ModelBase):
def on_build(self, patch_size, in_ch, base_ch = 16, use_fp16 = False): def on_build(self, patch_size, in_ch, base_ch = 16, use_fp16 = False):
self.use_fp16 = use_fp16 self.use_fp16 = use_fp16
conv_dtype = tf.float16 if use_fp16 else tf.float32 conv_dtype = tf.float16 if use_fp16 else tf.float32
class ResidualBlock(nn.ModelBase): class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3 ): def on_build(self, ch, kernel_size=3 ):
self.conv1 = 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=conv_dtype)
@ -150,7 +150,7 @@ class UNetPatchDiscriminator(nn.ModelBase):
self.convs = [] self.convs = []
self.upconvs = [] self.upconvs = []
layers = self.find_archi(patch_size) 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)]) } 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) self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID', dtype=conv_dtype)
@ -169,14 +169,14 @@ class UNetPatchDiscriminator(nn.ModelBase):
def forward(self, x): def forward(self, x):
if self.use_fp16: if self.use_fp16:
x = tf.cast(x, tf.float16) x = tf.cast(x, tf.float16)
x = tf.nn.leaky_relu( self.in_conv(x), 0.2 ) x = tf.nn.leaky_relu( self.in_conv(x), 0.2 )
encs = [] encs = []
for conv in self.convs: for conv in self.convs:
encs.insert(0, x) encs.insert(0, x)
x = tf.nn.leaky_relu( conv(x), 0.2 ) x = tf.nn.leaky_relu( conv(x), 0.2 )
center_out, x = self.center_out(x), tf.nn.leaky_relu( self.center_conv(x), 0.2 ) 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) in enumerate(zip(self.upconvs, encs)):
@ -184,7 +184,7 @@ class UNetPatchDiscriminator(nn.ModelBase):
x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis) x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis)
x = self.out_conv(x) x = self.out_conv(x)
if self.use_fp16: if self.use_fp16:
center_out = tf.cast(center_out, tf.float32) center_out = tf.cast(center_out, tf.float32)
x = tf.cast(x, tf.float32) x = tf.cast(x, tf.float32)

42
main.py
View file

@ -23,7 +23,7 @@ if __name__ == "__main__":
setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values))) setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))
exit_code = 0 exit_code = 0
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers() subparsers = parser.add_subparsers()
@ -52,9 +52,9 @@ if __name__ == "__main__":
p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to <output-dir>_debug\ directory.") p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to <output-dir>_debug\ directory.")
p.add_argument('--no-output-debug', action="store_false", dest="output_debug", default=None, help="Don't writes debug images to <output-dir>_debug\ directory.") p.add_argument('--no-output-debug', action="store_false", dest="output_debug", default=None, help="Don't writes debug images to <output-dir>_debug\ directory.")
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'whole_face', 'head', 'mark_only'], default=None) p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'whole_face', 'head', 'mark_only'], default=None)
p.add_argument('--max-faces-from-image', type=int, dest="max_faces_from_image", default=None, help="Max faces from image.") p.add_argument('--max-faces-from-image', type=int, dest="max_faces_from_image", default=None, help="Max faces from image.")
p.add_argument('--image-size', type=int, dest="image_size", default=None, help="Output image size.") p.add_argument('--image-size', type=int, dest="image_size", default=None, help="Output image size.")
p.add_argument('--jpeg-quality', type=int, dest="jpeg_quality", default=None, help="Jpeg quality.") p.add_argument('--jpeg-quality', type=int, dest="jpeg_quality", default=None, help="Jpeg quality.")
p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.") p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")
p.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.") p.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.")
p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.")
@ -144,10 +144,10 @@ if __name__ == "__main__":
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") 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('--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('--silent-start', action="store_true", dest="silent_start", default=False, help="Silent start. Automatically chooses Best GPU and last used model.")
p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+') p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+')
p.set_defaults (func=process_train) p.set_defaults (func=process_train)
def process_exportdfm(arguments): def process_exportdfm(arguments):
osex.set_process_lowest_prio() osex.set_process_lowest_prio()
from mainscripts import ExportDFM from mainscripts import ExportDFM
@ -261,8 +261,8 @@ if __name__ == "__main__":
p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.") p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
p.set_defaults(func=process_faceset_enhancer) p.set_defaults(func=process_faceset_enhancer)
p = facesettool_parser.add_parser ("resize", help="Resize DFL faceset.") p = facesettool_parser.add_parser ("resize", help="Resize DFL faceset.")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
@ -271,7 +271,7 @@ if __name__ == "__main__":
from mainscripts import FacesetResizer from mainscripts import FacesetResizer
FacesetResizer.process_folder ( Path(arguments.input_dir) ) FacesetResizer.process_folder ( Path(arguments.input_dir) )
p.set_defaults(func=process_faceset_resizer) p.set_defaults(func=process_faceset_resizer)
def process_dev_test(arguments): def process_dev_test(arguments):
osex.set_process_lowest_prio() osex.set_process_lowest_prio()
from mainscripts import dev_misc from mainscripts import dev_misc
@ -280,10 +280,10 @@ if __name__ == "__main__":
p = subparsers.add_parser( "dev_test", help="") p = subparsers.add_parser( "dev_test", help="")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_dev_test) p.set_defaults (func=process_dev_test)
# ========== XSeg # ========== XSeg
xseg_parser = subparsers.add_parser( "xseg", help="XSeg tools.").add_subparsers() xseg_parser = subparsers.add_parser( "xseg", help="XSeg tools.").add_subparsers()
p = xseg_parser.add_parser( "editor", help="XSeg editor.") p = xseg_parser.add_parser( "editor", help="XSeg editor.")
def process_xsegeditor(arguments): def process_xsegeditor(arguments):
@ -291,11 +291,11 @@ if __name__ == "__main__":
from XSegEditor import XSegEditor from XSegEditor import XSegEditor
global exit_code global exit_code
exit_code = XSegEditor.start (Path(arguments.input_dir)) exit_code = XSegEditor.start (Path(arguments.input_dir))
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_xsegeditor) p.set_defaults (func=process_xsegeditor)
p = xseg_parser.add_parser( "apply", help="Apply trained XSeg model to the extracted faces.") p = xseg_parser.add_parser( "apply", help="Apply trained XSeg model to the extracted faces.")
def process_xsegapply(arguments): def process_xsegapply(arguments):
@ -305,8 +305,8 @@ if __name__ == "__main__":
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir")
p.set_defaults (func=process_xsegapply) p.set_defaults (func=process_xsegapply)
p = xseg_parser.add_parser( "remove", help="Remove applied XSeg masks from the extracted faces.") p = xseg_parser.add_parser( "remove", help="Remove applied XSeg masks from the extracted faces.")
def process_xsegremove(arguments): def process_xsegremove(arguments):
osex.set_process_lowest_prio() osex.set_process_lowest_prio()
@ -314,8 +314,8 @@ if __name__ == "__main__":
XSegUtil.remove_xseg (Path(arguments.input_dir) ) XSegUtil.remove_xseg (Path(arguments.input_dir) )
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_xsegremove) p.set_defaults (func=process_xsegremove)
p = xseg_parser.add_parser( "remove_labels", help="Remove XSeg labels from the extracted faces.") p = xseg_parser.add_parser( "remove_labels", help="Remove XSeg labels from the extracted faces.")
def process_xsegremovelabels(arguments): def process_xsegremovelabels(arguments):
osex.set_process_lowest_prio() osex.set_process_lowest_prio()
@ -323,8 +323,8 @@ if __name__ == "__main__":
XSegUtil.remove_xseg_labels (Path(arguments.input_dir) ) XSegUtil.remove_xseg_labels (Path(arguments.input_dir) )
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_xsegremovelabels) p.set_defaults (func=process_xsegremovelabels)
p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in <input_dir>_xseg dir.") p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in <input_dir>_xseg dir.")
def process_xsegfetch(arguments): def process_xsegfetch(arguments):
@ -333,7 +333,7 @@ if __name__ == "__main__":
XSegUtil.fetch_xseg (Path(arguments.input_dir) ) XSegUtil.fetch_xseg (Path(arguments.input_dir) )
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
p.set_defaults (func=process_xsegfetch) p.set_defaults (func=process_xsegfetch)
def bad_args(arguments): def bad_args(arguments):
parser.print_help() parser.print_help()
exit(0) exit(0)
@ -344,9 +344,9 @@ if __name__ == "__main__":
if exit_code == 0: if exit_code == 0:
print ("Done.") print ("Done.")
exit(exit_code) exit(exit_code)
''' '''
import code import code
code.interact(local=dict(globals(), **locals())) code.interact(local=dict(globals(), **locals()))

2
mainscripts/ExportDFM.py
View file

@ -19,4 +19,4 @@ def main(model_class_name, saved_models_path):
is_exporting=True, is_exporting=True,
saved_models_path=saved_models_path, saved_models_path=saved_models_path,
cpu_only=True) cpu_only=True)
model.export_dfm () model.export_dfm ()

50
mainscripts/FacesetResizer.py
View file

@ -79,79 +79,79 @@ class FacesetResizerSubprocessor(Subprocessor):
h,w = img.shape[:2] h,w = img.shape[:2]
if h != w: if h != w:
raise Exception(f'w != h in {filepath}') raise Exception(f'w != h in {filepath}')
image_size = self.image_size image_size = self.image_size
face_type = self.face_type face_type = self.face_type
output_filepath = self.output_dirpath / filepath.name output_filepath = self.output_dirpath / filepath.name
if face_type is not None: if face_type is not None:
lmrks = dflimg.get_landmarks() lmrks = dflimg.get_landmarks()
mat = LandmarksProcessor.get_transform_mat(lmrks, image_size, face_type) mat = LandmarksProcessor.get_transform_mat(lmrks, image_size, face_type)
img = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 ) img = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 )
img = np.clip(img, 0, 255).astype(np.uint8) img = np.clip(img, 0, 255).astype(np.uint8)
cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
dfl_dict = dflimg.get_dict() dfl_dict = dflimg.get_dict()
dflimg = DFLIMG.load (output_filepath) dflimg = DFLIMG.load (output_filepath)
dflimg.set_dict(dfl_dict) dflimg.set_dict(dfl_dict)
xseg_mask = dflimg.get_xseg_mask() xseg_mask = dflimg.get_xseg_mask()
if xseg_mask is not None: if xseg_mask is not None:
xseg_res = 256 xseg_res = 256
xseg_lmrks = lmrks.copy() xseg_lmrks = lmrks.copy()
xseg_lmrks *= (xseg_res / w) xseg_lmrks *= (xseg_res / w)
xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type) xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type)
xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 ) xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 )
xseg_mask[xseg_mask < 0.5] = 0 xseg_mask[xseg_mask < 0.5] = 0
xseg_mask[xseg_mask >= 0.5] = 1 xseg_mask[xseg_mask >= 0.5] = 1
dflimg.set_xseg_mask(xseg_mask) dflimg.set_xseg_mask(xseg_mask)
seg_ie_polys = dflimg.get_seg_ie_polys() seg_ie_polys = dflimg.get_seg_ie_polys()
for poly in seg_ie_polys.get_polys(): for poly in seg_ie_polys.get_polys():
poly_pts = poly.get_pts() poly_pts = poly.get_pts()
poly_pts = LandmarksProcessor.transform_points(poly_pts, mat) poly_pts = LandmarksProcessor.transform_points(poly_pts, mat)
poly.set_points(poly_pts) poly.set_points(poly_pts)
dflimg.set_seg_ie_polys(seg_ie_polys) dflimg.set_seg_ie_polys(seg_ie_polys)
lmrks = LandmarksProcessor.transform_points(lmrks, mat) lmrks = LandmarksProcessor.transform_points(lmrks, mat)
dflimg.set_landmarks(lmrks) dflimg.set_landmarks(lmrks)
image_to_face_mat = dflimg.get_image_to_face_mat() image_to_face_mat = dflimg.get_image_to_face_mat()
if image_to_face_mat is not None: if image_to_face_mat is not None:
image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type ) image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
dflimg.set_image_to_face_mat(image_to_face_mat) dflimg.set_image_to_face_mat(image_to_face_mat)
dflimg.set_face_type( FaceType.toString(face_type) ) dflimg.set_face_type( FaceType.toString(face_type) )
dflimg.save() dflimg.save()
else: else:
dfl_dict = dflimg.get_dict() dfl_dict = dflimg.get_dict()
scale = w / image_size scale = w / image_size
img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4) img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4)
cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
dflimg = DFLIMG.load (output_filepath) dflimg = DFLIMG.load (output_filepath)
dflimg.set_dict(dfl_dict) dflimg.set_dict(dfl_dict)
lmrks = dflimg.get_landmarks() lmrks = dflimg.get_landmarks()
lmrks /= scale lmrks /= scale
dflimg.set_landmarks(lmrks) dflimg.set_landmarks(lmrks)
seg_ie_polys = dflimg.get_seg_ie_polys() seg_ie_polys = dflimg.get_seg_ie_polys()
seg_ie_polys.mult_points( 1.0 / scale) seg_ie_polys.mult_points( 1.0 / scale)
dflimg.set_seg_ie_polys(seg_ie_polys) dflimg.set_seg_ie_polys(seg_ie_polys)
image_to_face_mat = dflimg.get_image_to_face_mat() image_to_face_mat = dflimg.get_image_to_face_mat()
if image_to_face_mat is not None: if image_to_face_mat is not None:
face_type = FaceType.fromString ( dflimg.get_face_type() ) face_type = FaceType.fromString ( dflimg.get_face_type() )
image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type ) image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
@ -165,9 +165,9 @@ class FacesetResizerSubprocessor(Subprocessor):
return (0, filepath, None) return (0, filepath, None)
def process_folder ( dirpath): def process_folder ( dirpath):
image_size = io.input_int(f"New image size", 512, valid_range=[128,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() face_type = io.input_str ("Change face type", 'same', ['h','mf','f','wf','head','same']).lower()
if face_type == 'same': if face_type == 'same':
face_type = None face_type = None
@ -177,7 +177,7 @@ def process_folder ( dirpath):
'f' : FaceType.FULL, 'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE, 'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[face_type] 'head' : FaceType.HEAD}[face_type]
output_dirpath = dirpath.parent / (dirpath.name + '_resized') output_dirpath = dirpath.parent / (dirpath.name + '_resized')
output_dirpath.mkdir (exist_ok=True, parents=True) output_dirpath.mkdir (exist_ok=True, parents=True)

10
mainscripts/Trainer.py
View file

@ -42,7 +42,7 @@ def trainerThread (s2c, c2s, e,
if not saved_models_path.exists(): if not saved_models_path.exists():
saved_models_path.mkdir(exist_ok=True, parents=True) saved_models_path.mkdir(exist_ok=True, parents=True)
model = models.import_model(model_class_name)( model = models.import_model(model_class_name)(
is_training=True, is_training=True,
saved_models_path=saved_models_path, saved_models_path=saved_models_path,
@ -67,10 +67,10 @@ def trainerThread (s2c, c2s, e,
io.log_info ("Saving....", end='\r') io.log_info ("Saving....", end='\r')
model.save() model.save()
shared_state['after_save'] = True shared_state['after_save'] = True
def model_backup(): def model_backup():
if not debug and not is_reached_goal: if not debug and not is_reached_goal:
model.create_backup() model.create_backup()
def send_preview(): def send_preview():
if not debug: if not debug:
@ -119,7 +119,7 @@ def trainerThread (s2c, c2s, e,
io.log_info("") io.log_info("")
io.log_info("Trying to do the first iteration. If an error occurs, reduce the model parameters.") io.log_info("Trying to do the first iteration. If an error occurs, reduce the model parameters.")
io.log_info("") io.log_info("")
if sys.platform[0:3] == 'win': if sys.platform[0:3] == 'win':
io.log_info("!!!") io.log_info("!!!")
io.log_info("Windows 10 users IMPORTANT notice. You should set this setting in order to work correctly.") io.log_info("Windows 10 users IMPORTANT notice. You should set this setting in order to work correctly.")
@ -137,7 +137,7 @@ def trainerThread (s2c, c2s, e,
if shared_state['after_save']: if shared_state['after_save']:
shared_state['after_save'] = False shared_state['after_save'] = False
mean_loss = np.mean ( loss_history[save_iter:iter], axis=0) mean_loss = np.mean ( loss_history[save_iter:iter], axis=0)
for loss_value in mean_loss: for loss_value in mean_loss:

80
models/Model_SAEHD/Model.py
View file

@ -29,7 +29,7 @@ class SAEHDModel(ModelBase):
yn_str = {True:'y',False:'n'} yn_str = {True:'y',False:'n'}
min_res = 64 min_res = 64
max_res = 640 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_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_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'f')
@ -70,14 +70,14 @@ class SAEHDModel(ModelBase):
self.ask_random_dst_flip() self.ask_random_dst_flip()
self.ask_batch_size(suggest_batch_size) 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(): 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 = 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) resolution = np.clip ( (resolution // 16) * 16, min_res, max_res)
self.options['resolution'] = resolution self.options['resolution'] = resolution
self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head'], help_message="Half / mid face / full face / whole face / head. Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset.").lower() self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head'], help_message="Half / mid face / full face / whole face / head. Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset.").lower()
while True: while True:
@ -138,11 +138,11 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
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['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.') 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) default_gan_power = self.options['gan_power'] = self.load_or_def_option('gan_power', 0.0)
default_gan_patch_size = self.options['gan_patch_size'] = self.load_or_def_option('gan_patch_size', self.options['resolution'] // 8) 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_dims = self.options['gan_dims'] = self.load_or_def_option('gan_dims', 16)
if self.is_first_run() or ask_override: 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['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.")
@ -153,14 +153,14 @@ 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['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 .. 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 ) 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: 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 ) 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 self.options['gan_patch_size'] = gan_patch_size
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 ) 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['gan_dims'] = gan_dims
if 'df' in self.options['archi']: if 'df' in self.options['archi']:
self.options['true_face_power'] = np.clip ( io.input_number ("'True face' power.", default_true_face_power, add_info="0.0000 .. 1.0", help_message="Experimental option. Discriminates result face to be more like src face. Higher value - stronger discrimination. Typical value is 0.01 . Comparison - https://i.imgur.com/czScS9q.png"), 0.0, 1.0 ) self.options['true_face_power'] = np.clip ( io.input_number ("'True face' power.", default_true_face_power, add_info="0.0000 .. 1.0", help_message="Experimental option. Discriminates result face to be more like src face. Higher value - stronger discrimination. Typical value is 0.01 . Comparison - https://i.imgur.com/czScS9q.png"), 0.0, 1.0 )
else: else:
@ -176,7 +176,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
if self.options['pretrain'] and self.get_pretraining_data_path() is None: if self.options['pretrain'] and self.get_pretraining_data_path() is None:
raise Exception("pretraining_data_path is not defined") raise Exception("pretraining_data_path is not defined")
self.gan_model_changed = (default_gan_patch_size != self.options['gan_patch_size']) or (default_gan_dims != self.options['gan_dims']) 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)
@ -198,7 +198,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
if 'eyes_prio' in self.options: if 'eyes_prio' in self.options:
self.options.pop('eyes_prio') self.options.pop('eyes_prio')
eyes_mouth_prio = self.options['eyes_mouth_prio'] eyes_mouth_prio = self.options['eyes_mouth_prio']
archi_split = self.options['archi'].split('-') archi_split = self.options['archi'].split('-')
@ -207,7 +207,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
archi_type, archi_opts = archi_split archi_type, archi_opts = archi_split
elif len(archi_split) == 1: elif len(archi_split) == 1:
archi_type, archi_opts = archi_split[0], None archi_type, archi_opts = archi_split[0], None
self.archi_type = archi_type self.archi_type = archi_type
ae_dims = self.options['ae_dims'] ae_dims = self.options['ae_dims']
@ -220,12 +220,12 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
adabelief = self.options['adabelief'] adabelief = self.options['adabelief']
use_fp16 = False#self.options['use_fp16'] use_fp16 = False#self.options['use_fp16']
self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power'] 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_warp = False if self.pretrain else self.options['random_warp']
random_src_flip = self.random_src_flip if not self.pretrain else True 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 random_dst_flip = self.random_dst_flip if not self.pretrain else True
if self.pretrain: if self.pretrain:
self.options_show_override['gan_power'] = 0.0 self.options_show_override['gan_power'] = 0.0
self.options_show_override['random_warp'] = False self.options_show_override['random_warp'] = False
@ -238,8 +238,8 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
ct_mode = self.options['ct_mode'] ct_mode = self.options['ct_mode']
if ct_mode == 'none': if ct_mode == 'none':
ct_mode = None ct_mode = None
models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu'] 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' 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' optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
@ -353,7 +353,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
gpu_G_loss_gvs = [] gpu_G_loss_gvs = []
gpu_D_code_loss_gvs = [] gpu_D_code_loss_gvs = []
gpu_D_src_dst_loss_gvs = [] gpu_D_src_dst_loss_gvs = []
for gpu_id in range(gpu_count): 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' ): with tf.device( f'/{devices[gpu_id].tf_dev_type}:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
with tf.device(f'/CPU:0'): with tf.device(f'/CPU:0'):
@ -405,7 +405,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
gpu_target_dstm_style_blur = gpu_target_dstm_blur #default style mask is 0.5 on boundary gpu_target_dstm_style_blur = gpu_target_dstm_blur #default style mask is 0.5 on boundary
gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_blur, 0, 0.5) * 2 gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_blur, 0, 0.5) * 2
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
gpu_target_dst_style_masked = gpu_target_dst*gpu_target_dstm_style_blur gpu_target_dst_style_masked = gpu_target_dst*gpu_target_dstm_style_blur
gpu_target_dst_style_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_style_blur) gpu_target_dst_style_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_style_blur)
@ -500,14 +500,14 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
gpu_G_loss += gan_power*(DLoss(gpu_pred_src_src_d_ones, gpu_pred_src_src_d) + \ gpu_G_loss += gan_power*(DLoss(gpu_pred_src_src_d_ones, gpu_pred_src_src_d) + \
DLoss(gpu_pred_src_src_d2_ones, gpu_pred_src_src_d2)) DLoss(gpu_pred_src_src_d2_ones, gpu_pred_src_src_d2))
if masked_training: if masked_training:
# Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan # 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.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 += 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.src_dst_trainable_weights ) ] gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
@ -617,10 +617,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
if do_init: if do_init:
model.init_weights() model.init_weights()
############### ###############
# initializing sample generators # initializing sample generators
if self.is_training: 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_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
@ -661,20 +661,20 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
if self.pretrain_just_disabled: if self.pretrain_just_disabled:
self.update_sample_for_preview(force_new=True) self.update_sample_for_preview(force_new=True)
def export_dfm (self): def export_dfm (self):
output_path=self.get_strpath_storage_for_file('model.dfm') output_path=self.get_strpath_storage_for_file('model.dfm')
io.log_info(f'Dumping .dfm to {output_path}') io.log_info(f'Dumping .dfm to {output_path}')
tf = nn.tf tf = nn.tf
nn.set_data_format('NCHW') nn.set_data_format('NCHW')
with tf.device (nn.tf_default_device_name): 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.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face')
warped_dst = tf.transpose(warped_dst, (0,3,1,2)) warped_dst = tf.transpose(warped_dst, (0,3,1,2))
if 'df' in self.archi_type: if 'df' in self.archi_type:
gpu_dst_code = self.inter(self.encoder(warped_dst)) gpu_dst_code = self.inter(self.encoder(warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
@ -689,21 +689,21 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code) gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_code) _, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
gpu_pred_src_dst = tf.transpose(gpu_pred_src_dst, (0,2,3,1)) gpu_pred_src_dst = tf.transpose(gpu_pred_src_dst, (0,2,3,1))
gpu_pred_dst_dstm = tf.transpose(gpu_pred_dst_dstm, (0,2,3,1)) 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)) gpu_pred_src_dstm = tf.transpose(gpu_pred_src_dstm, (0,2,3,1))
tf.identity(gpu_pred_dst_dstm, name='out_face_mask') 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_dst, name='out_celeb_face')
tf.identity(gpu_pred_src_dstm, name='out_celeb_face_mask') tf.identity(gpu_pred_src_dstm, name='out_celeb_face_mask')
output_graph_def = tf.graph_util.convert_variables_to_constants( output_graph_def = tf.graph_util.convert_variables_to_constants(
nn.tf_sess, nn.tf_sess,
tf.get_default_graph().as_graph_def(), tf.get_default_graph().as_graph_def(),
['out_face_mask','out_celeb_face','out_celeb_face_mask'] ['out_face_mask','out_celeb_face','out_celeb_face_mask']
) )
import tf2onnx import tf2onnx
with tf.device("/CPU:0"): with tf.device("/CPU:0"):
model_proto, _ = tf2onnx.convert._convert_common( model_proto, _ = tf2onnx.convert._convert_common(
@ -713,7 +713,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
output_names=['out_face_mask:0','out_celeb_face:0','out_celeb_face_mask:0'], output_names=['out_face_mask:0','out_celeb_face:0','out_celeb_face_mask:0'],
opset=13, opset=13,
output_path=output_path) output_path=output_path)
#override #override
def get_model_filename_list(self): def get_model_filename_list(self):
return self.model_filename_list return self.model_filename_list