1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-09 04:36:50 -04:00
faceswap/lib/cli/launcher.py
torzdf d8557c1970
Faceswap 2.0 (#1045)
* Core Updates
    - Remove lib.utils.keras_backend_quiet and replace with get_backend() where relevant
    - Document lib.gpu_stats and lib.sys_info
    - Remove call to GPUStats.is_plaidml from convert and replace with get_backend()
    - lib.gui.menu - typofix

* Update Dependencies
Bump Tensorflow Version Check

* Port extraction to tf2

* Add custom import finder for loading Keras or tf.keras depending on backend

* Add `tensorflow` to KerasFinder search path

* Basic TF2 training running

* model.initializers - docstring fix

* Fix and pass tests for tf2

* Replace Keras backend tests with faceswap backend tests

* Initial optimizers update

* Monkey patch tf.keras optimizer

* Remove custom Adam Optimizers and Memory Saving Gradients

* Remove multi-gpu option. Add Distribution to cli

* plugins.train.model._base: Add Mirror, Central and Default distribution strategies

* Update tensorboard kwargs for tf2

* Penalized Loss - Fix for TF2 and AMD

* Fix syntax for tf2.1

* requirements typo fix

* Explicit None for clipnorm if using a distribution strategy

* Fix penalized loss for distribution strategies

* Update Dlight

* typo fix

* Pin to TF2.2

* setup.py - Install tensorflow from pip if not available in Conda

* Add reduction options and set default for mirrored distribution strategy

* Explicitly use default strategy rather than nullcontext

* lib.model.backup_restore documentation

* Remove mirrored strategy reduction method and default based on OS

* Initial restructure - training

* Remove PingPong
Start model.base refactor

* Model saving and resuming enabled

* More tidying up of model.base

* Enable backup and snapshotting

* Re-enable state file
Remove loss names from state file
Fix print loss function
Set snapshot iterations correctly

* Revert original model to Keras Model structure rather than custom layer
Output full model and sub model summary
Change NNBlocks to callables rather than custom keras layers

* Apply custom Conv2D layer

* Finalize NNBlock restructure
Update Dfaker blocks

* Fix reloading model under a different distribution strategy

* Pass command line arguments through to trainer

* Remove training_opts from model and reference params directly

* Tidy up model __init__

* Re-enable tensorboard logging
Suppress "Model Not Compiled" warning

* Fix timelapse

* lib.model.nnblocks - Bugfix residual block
Port dfaker
bugfix original

* dfl-h128 ported

* DFL SAE ported

* IAE Ported

* dlight ported

* port lightweight

* realface ported

* unbalanced ported

* villain ported

* lib.cli.args - Update Batchsize + move allow_growth to config

* Remove output shape definition
Get image sizes per side rather than globally

* Strip mask input from encoder

* Fix learn mask and output learned mask to preview

* Trigger Allow Growth prior to setting strategy

* Fix GUI Graphing

* GUI - Display batchsize correctly + fix training graphs

* Fix penalized loss

* Enable mixed precision training

* Update analysis displayed batch to match input

* Penalized Loss - Multi-GPU Fix

* Fix all losses for TF2

* Fix Reflect Padding

* Allow different input size for each side of the model

* Fix conv-aware initialization on reload

* Switch allow_growth order

* Move mixed_precision to cli

* Remove distrubution strategies

* Compile penalized loss sub-function into LossContainer

* Bump default save interval to 250
Generate preview on first iteration but don't save
Fix iterations to start at 1 instead of 0
Remove training deprecation warnings
Bump some scripts.train loglevels

* Add ability to refresh preview on demand on pop-up window

* Enable refresh of training preview from GUI

* Fix Convert
Debug logging in Initializers

* Fix Preview Tool

* Update Legacy TF1 weights to TF2
Catch stats error on loading stats with missing logs

* lib.gui.popup_configure - Make more responsive + document

* Multiple Outputs supported in trainer
Original Model - Mask output bugfix

* Make universal inference model for convert
Remove scaling from penalized mask loss (now handled at input to y_true)

* Fix inference model to work properly with all models

* Fix multi-scale output for convert

* Fix clipnorm issue with distribution strategies
Edit error message on OOM

* Update plaidml losses

* Add missing file

* Disable gmsd loss for plaidnl

* PlaidML - Basic training working

* clipnorm rewriting for mixed-precision

* Inference model creation bugfixes

* Remove debug code

* Bugfix: Default clipnorm to 1.0

* Remove all mask inputs from training code

* Remove mask inputs from convert

* GUI - Analysis Tab - Docstrings

* Fix rate in totals row

* lib.gui - Only update display pages if they have focus

* Save the model on first iteration

* plaidml - Fix SSIM loss with penalized loss

* tools.alignments - Remove manual and fix jobs

* GUI - Remove case formatting on help text

* gui MultiSelect custom widget - Set default values on init

* vgg_face2 - Move to plugins.extract.recognition and use plugins._base base class
cli - Add global GPU Exclude Option
tools.sort - Use global GPU Exlude option for backend
lib.model.session - Exclude all GPUs when running in CPU mode
lib.cli.launcher - Set backend to CPU mode when all GPUs excluded

* Cascade excluded devices to GPU Stats

* Explicit GPU selection for Train and Convert

* Reduce Tensorflow Min GPU Multiprocessor Count to 4

* remove compat.v1 code from extract

* Force TF to skip mixed precision compatibility check if GPUs have been filtered

* Add notes to config for non-working AMD losses

* Rasie error if forcing extract to CPU mode

* Fix loading of legace dfl-sae weights + dfl-sae typo fix

* Remove unused requirements
Update sphinx requirements
Fix broken rst file locations

* docs: lib.gui.display

* clipnorm amd condition check

* documentation - gui.display_analysis

* Documentation - gui.popup_configure

* Documentation - lib.logger

* Documentation - lib.model.initializers

* Documentation - lib.model.layers

* Documentation - lib.model.losses

* Documentation - lib.model.nn_blocks

* Documetation - lib.model.normalization

* Documentation - lib.model.session

* Documentation - lib.plaidml_stats

* Documentation: lib.training_data

* Documentation: lib.utils

* Documentation: plugins.train.model._base

* GUI Stats: prevent stats from using GPU

* Documentation - Original Model

* Documentation: plugins.model.trainer._base

* linting

* unit tests: initializers + losses

* unit tests: nn_blocks

* bugfix - Exclude gpu devices in train, not include

* Enable Exclude-Gpus in Extract

* Enable exclude gpus in tools

* Disallow multiple plugin types in a single model folder

* Automatically add exclude_gpus argument in for cpu backends

* Cpu backend fixes

* Relax optimizer test threshold

* Default Train settings - Set mask to Extended

* Update Extractor cli help text
Update to Python 3.8

* Fix FAN to run on CPU

* lib.plaidml_tools - typofix

* Linux installer - check for curl

* linux installer - typo fix
2020-08-12 10:36:41 +01:00

246 lines
10 KiB
Python

#!/usr/bin/env python3
""" Launches the correct script with the given Command Line Arguments """
import logging
import os
import platform
import sys
from importlib import import_module
from lib.gpu_stats import set_exclude_devices, GPUStats
from lib.logger import crash_log, log_setup
from lib.utils import (FaceswapError, get_backend, KerasFinder, safe_shutdown, set_backend,
set_system_verbosity)
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
class ScriptExecutor(): # pylint:disable=too-few-public-methods
""" Loads the relevant script modules and executes the script.
This class is initialized in each of the argparsers for the relevant
command, then execute script is called within their set_default
function.
Parameters
----------
command: str
The faceswap command that is being executed
"""
def __init__(self, command):
self._command = command.lower()
def _import_script(self):
""" Imports the relevant script as indicated by :attr:`_command` from the scripts folder.
Returns
-------
class: Faceswap Script
The uninitialized script from the faceswap scripts folder.
"""
self._test_for_tf_version()
self._test_for_gui()
cmd = os.path.basename(sys.argv[0])
src = "tools.{}".format(self._command.lower()) if cmd == "tools.py" else "scripts"
mod = ".".join((src, self._command.lower()))
module = import_module(mod)
script = getattr(module, self._command.title())
return script
@staticmethod
def _test_for_tf_version():
""" Check that the required Tensorflow version is installed.
Raises
------
FaceswapError
If Tensorflow is not found, or is not between versions 2.2 and 2.2
"""
min_ver = 2.2
max_ver = 2.2
try:
# Ensure tensorflow doesn't pin all threads to one core when using Math Kernel Library
os.environ["TF_MIN_GPU_MULTIPROCESSOR_COUNT"] = "4"
os.environ["KMP_AFFINITY"] = "disabled"
import tensorflow as tf # pylint:disable=import-outside-toplevel
except ImportError as err:
raise FaceswapError("There was an error importing Tensorflow. This is most likely "
"because you do not have TensorFlow installed, or you are trying "
"to run tensorflow-gpu on a system without an Nvidia graphics "
"card. Original import error: {}".format(str(err)))
tf_ver = float(".".join(tf.__version__.split(".")[:2])) # pylint:disable=no-member
if tf_ver < min_ver:
raise FaceswapError("The minimum supported Tensorflow is version {} but you have "
"version {} installed. Please upgrade Tensorflow.".format(
min_ver, tf_ver))
if tf_ver > max_ver:
raise FaceswapError("The maximumum supported Tensorflow is version {} but you have "
"version {} installed. Please downgrade Tensorflow.".format(
max_ver, tf_ver))
logger.debug("Installed Tensorflow Version: %s", tf_ver)
def _test_for_gui(self):
""" If running the gui, performs check to ensure necessary prerequisites are present. """
if self._command != "gui":
return
self._test_tkinter()
self._check_display()
@staticmethod
def _test_tkinter():
""" If the user is running the GUI, test whether the tkinter app is available on their
machine. If not exit gracefully.
This avoids having to import every tkinter function within the GUI in a wrapper and
potentially spamming traceback errors to console.
Raises
------
FaceswapError
If tkinter cannot be imported
"""
try:
# pylint: disable=unused-variable
import tkinter # noqa pylint: disable=unused-import,import-outside-toplevel
except ImportError:
logger.error("It looks like TkInter isn't installed for your OS, so the GUI has been "
"disabled. To enable the GUI please install the TkInter application. You "
"can try:")
logger.info("Anaconda: conda install tk")
logger.info("Windows/macOS: Install ActiveTcl Community Edition from "
"http://www.activestate.com")
logger.info("Ubuntu/Mint/Debian: sudo apt install python3-tk")
logger.info("Arch: sudo pacman -S tk")
logger.info("CentOS/Redhat: sudo yum install tkinter")
logger.info("Fedora: sudo dnf install python3-tkinter")
raise FaceswapError("TkInter not found")
@staticmethod
def _check_display():
""" Check whether there is a display to output the GUI to.
If running on Windows then it is assumed that we are not running in headless mode
Raises
------
FaceswapError
If a DISPLAY environmental cannot be found
"""
if not os.environ.get("DISPLAY", None) and os.name != "nt":
if platform.system() == "Darwin":
logger.info("macOS users need to install XQuartz. "
"See https://support.apple.com/en-gb/HT201341")
raise FaceswapError("No display detected. GUI mode has been disabled.")
def execute_script(self, arguments):
""" Performs final set up and launches the requested :attr:`_command` with the given
command line arguments.
Monitors for errors and attempts to shut down the process cleanly on exit.
Parameters
----------
arguments: :class:`argparse.Namespace`
The command line arguments to be passed to the executing script.
"""
set_system_verbosity(arguments.loglevel)
is_gui = hasattr(arguments, "redirect_gui") and arguments.redirect_gui
log_setup(arguments.loglevel, arguments.logfile, self._command, is_gui)
success = False
if self._command != "gui":
self._configure_backend(arguments)
try:
script = self._import_script()
process = script(arguments)
process.process()
success = True
except FaceswapError as err:
for line in str(err).splitlines():
logger.error(line)
except KeyboardInterrupt: # pylint: disable=try-except-raise
raise
except SystemExit:
pass
except Exception: # pylint: disable=broad-except
crash_file = crash_log()
logger.exception("Got Exception on main handler:")
logger.critical("An unexpected crash has occurred. Crash report written to '%s'. "
"You MUST provide this file if seeking assistance. Please verify you "
"are running the latest version of faceswap before reporting",
crash_file)
finally:
safe_shutdown(got_error=not success)
def _configure_backend(self, arguments):
""" Configure the backend.
Exclude any GPUs for use by Faceswap when requested.
Set Faceswap backend to CPU if all GPUs have been deselected.
Add the Keras import interception code.
Parameters
----------
arguments: :class:`argparse.Namespace`
The command line arguments passed to Faceswap.
"""
if not hasattr(arguments, "exclude_gpus"):
# Cpu backends will not have this attribute
logger.debug("Adding missing exclude gpus argument to namespace")
setattr(arguments, "exclude_gpus", None)
if arguments.exclude_gpus:
if not all(idx.isdigit() for idx in arguments.exclude_gpus):
logger.error("GPUs passed to the ['-X', '--exclude-gpus'] argument must all be "
"integers.")
sys.exit(1)
arguments.exclude_gpus = [int(idx) for idx in arguments.exclude_gpus]
set_exclude_devices(arguments.exclude_gpus)
if ((get_backend() == "cpu" or GPUStats().exclude_all_devices) and
(self._command == "extract" and arguments.detector in ("mtcnn", "s3fd"))):
logger.error("Extracting on CPU is not currently for detector: '%s'",
arguments.detector.upper())
sys.exit(0)
if GPUStats().exclude_all_devices and get_backend() != "cpu":
msg = "Switching backend to CPU"
if get_backend() == "amd":
msg += (". Using Tensorflow for CPU operations.")
os.environ["KERAS_BACKEND"] = "tensorflow"
set_backend("cpu")
logger.info(msg)
# Add Keras finder to the meta_path list as the first item
sys.meta_path.insert(0, KerasFinder())
logger.debug("Executing: %s. PID: %s", self._command, os.getpid())
if get_backend() == "amd":
plaidml_found = self._setup_amd(arguments)
if not plaidml_found:
safe_shutdown(got_error=True)
sys.exit(1)
@classmethod
def _setup_amd(cls, arguments):
""" Test for plaidml and perform setup for AMD.
Parameters
----------
arguments: :class:`argparse.Namespace`
The command line arguments passed to Faceswap.
"""
logger.debug("Setting up for AMD")
try:
import plaidml # noqa pylint:disable=unused-import,import-outside-toplevel
except ImportError:
logger.error("PlaidML not found. Run `pip install plaidml-keras` for AMD support")
return False
from lib.plaidml_tools import setup_plaidml # pylint:disable=import-outside-toplevel
setup_plaidml(arguments.loglevel, arguments.exclude_gpus)
logger.debug("setup up for PlaidML")
return True