Merge branch 'merge-upstream-changes' into patch-1

_requirements_base.txt

@@ -1,19 +1,16 @@
-tqdm>=4.62
+tqdm>=4.64
 psutil>=5.8.0
-numpy>=1.18.0,<1.20.0
-opencv-python>=4.5.3.0
+numpy>=1.18.0
+opencv-python>=4.5.5.0
 pillow>=8.3.1
-scikit-learn>=0.24.2
-fastcluster>=1.1.26
+scikit-learn>=1.0.2
+fastcluster>=1.2.4
 # matplotlib 3.3.1 breaks custom toolbar in graph popup
 matplotlib>=3.2.0,<3.3.0
 imageio>=2.9.0
-imageio-ffmpeg>=0.4.5
+imageio-ffmpeg>=0.4.7
 ffmpy==0.2.3
 # Exclude badly numbered Python2 version of nvidia-ml-py
-# nvidia-ml-py>=11.450,<300
-# v11.515.0 changes dtype of output items. Pinned for now
-# TODO update code to use latest version
-nvidia-ml-py>=11.450,<11.515
+nvidia-ml-py>=11.510,<300
 pywin32>=228 ; sys_platform == "win32"
 pynvx==1.0.0 ; sys_platform == "darwin"

lib/cli/launcher.py

@@ -9,8 +9,8 @@ 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)
+from lib.utils import (FaceswapError, get_backend, get_tf_version, safe_shutdown,
+                       set_backend, set_system_verbosity)
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
@@ -41,7 +41,7 @@ class ScriptExecutor():  # pylint:disable=too-few-public-methods
         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"
+        src = f"tools.{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())
@@ -53,15 +53,15 @@ class ScriptExecutor():  # pylint:disable=too-few-public-methods
         Raises
         ------
         FaceswapError
-            If Tensorflow is not found, or is not between versions 2.2 and 2.6
+            If Tensorflow is not found, or is not between versions 2.2 and 2.8
         """
         min_ver = 2.2
-        max_ver = 2.8 #2.6
+        max_ver = 2.8
         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
+            import tensorflow as tf  # noqa pylint:disable=import-outside-toplevel,unused-import
         except ImportError as err:
             if "DLL load failed while importing" in str(err):
                 msg = (
@@ -77,14 +77,14 @@ class ScriptExecutor():  # pylint:disable=too-few-public-methods
                     f"error: {str(err)}")
             self._handle_import_error(msg)
 
-        tf_ver = float(".".join(tf.__version__.split(".")[:2]))  # pylint:disable=no-member
+        tf_ver = get_tf_version()
         if tf_ver < min_ver:
-            msg = ("The minimum supported Tensorflow is version {} but you have version {} "
-                   "installed. Please upgrade Tensorflow.".format(min_ver, tf_ver))
+            msg = (f"The minimum supported Tensorflow is version {min_ver} but you have version "
+                   f"{tf_ver} installed. Please upgrade Tensorflow.")
             self._handle_import_error(msg)
         if tf_ver > max_ver:
-            msg = ("The maximum supported Tensorflow is version {} but you have version {} "
-                   "installed. Please downgrade Tensorflow.".format(max_ver, tf_ver))
+            msg = (f"The maximum supported Tensorflow is version {max_ver} but you have version "
+                   f"{tf_ver} installed. Please downgrade Tensorflow.")
             self._handle_import_error(msg)
         logger.debug("Installed Tensorflow Version: %s", tf_ver)
 
@@ -206,8 +206,6 @@ class ScriptExecutor():  # pylint:disable=too-few-public-methods
 
         Set Faceswap backend to CPU if all GPUs have been deselected.
 
-        Add the Keras import interception code.
-
         Parameters
         ----------
         arguments: :class:`argparse.Namespace`
@@ -234,9 +232,6 @@ class ScriptExecutor():  # pylint:disable=too-few-public-methods
             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":

lib/gui/analysis/event_reader.py

@@ -7,10 +7,12 @@ import zlib
 
 import numpy as np
 import tensorflow as tf
-from tensorflow.core.util import event_pb2
-from tensorflow.python.framework import errors_impl as tf_errors
+from tensorflow.core.util import event_pb2  # pylint:disable=no-name-in-module
+from tensorflow.python.framework import (  # pylint:disable=no-name-in-module
+    errors_impl as tf_errors)
 
 from lib.serializer import get_serializer
+from lib.utils import get_backend
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
@@ -43,7 +45,7 @@ class _LogFiles():
             The full path of each log file for each training session id that has been run
         """
         logger.debug("Loading log filenames. base_dir: '%s'", self._logs_folder)
-        retval = dict()
+        retval = {}
         for dirpath, _, filenames in os.walk(self._logs_folder):
             if not any(filename.startswith("events.out.tfevents") for filename in filenames):
                 continue
@@ -133,7 +135,7 @@ class _Cache():
     def __init__(self, session_ids):
         logger.debug("Initializing: %s: (session_ids: %s)", self.__class__.__name__, session_ids)
         self._data = {idx: None for idx in session_ids}
-        self._carry_over = dict()
+        self._carry_over = {}
         self._loss_labels = []
         logger.debug("Initialized: %s", self.__class__.__name__)
 
@@ -158,18 +160,19 @@ class _Cache():
         """
         logger.debug("Caching event data: (session_id: %s, labels: %s, data points: %s, "
                      "is_live: %s)", session_id, labels, len(data), is_live)
-        if not data:
-            logger.debug("No data to cache")
-            return
 
         if labels:
             logger.debug("Setting loss labels: %s", labels)
             self._loss_labels = labels
 
+        if not data:
+            logger.debug("No data to cache")
+            return
+
         timestamps, loss = self._to_numpy(data, is_live)
 
         if not is_live or (is_live and not self._data.get(session_id, None)):
-            self._data[session_id] = dict(labels=labels,
+            self._data[session_id] = dict(labels=self._loss_labels,
                                           loss=zlib.compress(loss),
                                           loss_shape=loss.shape,
                                           timestamps=zlib.compress(timestamps),
@@ -207,10 +210,30 @@ class _Cache():
                             for idx in sorted(data)])
         times, loss = self._process_data(data, times, loss, is_live)
 
-        times, loss = (np.array(times, dtype="float64"), np.array(loss, dtype="float32"))
+        if is_live and not all(len(val) == len(self._loss_labels) for val in loss):
+            # TODO Many attempts have been made to fix this for live graph logging, and the issue
+            # of non-consistent loss record sizes keeps coming up. In the meantime we shall swallow
+            # any loss values that are of incorrect length so graph remains functional. This will,
+            # most likely, lead to a mismatch on iteration count so a proper fix should be
+            # implemented.
 
+            # Timestamps and loss appears to remain consistent with each other, but sometimes loss
+            # appears non-consistent. eg (lengths):
+            # [2, 2, 2, 2, 2, 2, 2, 0] - last loss collection has zero length
+            # [1, 2, 2, 2, 2, 2, 2, 2] - 1st loss collection has 1 length
+            # [2, 2, 2, 3, 2, 2, 2] - 4th loss collection has 3 length
+
+            logger.debug("Inconsistent loss found in collection: %s", loss)
+            for idx in reversed(range(len(loss))):
+                if len(loss[idx]) != len(self._loss_labels):
+                    logger.debug("Removing loss/timestamps at position %s", idx)
+                    del loss[idx]
+                    del times[idx]
+
+        times, loss = (np.array(times, dtype="float64"), np.array(loss, dtype="float32"))
         logger.debug("Converted to numpy: (data points: %s, timestamps shape: %s, loss shape: %s)",
                      len(data), times.shape, loss.shape)
+
         return times, loss
 
     def _collect_carry_over(self, data):
@@ -334,7 +357,7 @@ class _Cache():
 
         dtype = "float32" if metric == "loss" else "float64"
 
-        retval = dict()
+        retval = {}
         for idx, data in raw.items():
             val = {metric: np.frombuffer(zlib.decompress(data[metric]),
                                          dtype=dtype).reshape(data[f"{metric}_shape"])}
@@ -461,7 +484,7 @@ class TensorBoardLogs():
             and list of loss values for each step
         """
         logger.debug("Getting loss: (session_id: %s)", session_id)
-        retval = dict()
+        retval = {}
         for idx in [session_id] if session_id else self.session_ids:
             self._check_cache(idx)
             data = self._cache.get_data(idx, "loss")
@@ -493,7 +516,7 @@ class TensorBoardLogs():
 
         logger.debug("Getting timestamps: (session_id: %s, is_training: %s)",
                      session_id, self._is_training)
-        retval = dict()
+        retval = {}
         for idx in [session_id] if session_id else self.session_ids:
             self._check_cache(idx)
             data = self._cache.get_data(idx, "timestamps")
@@ -565,7 +588,7 @@ class _EventParser():  # pylint:disable=too-few-public-methods
         session_id: int
             The session id that the data is being cached for
         """
-        data = dict()
+        data = {}
         try:
             for record in self._iterator:
                 event = event_pb2.Event.FromString(record)  # pylint:disable=no-member
@@ -573,8 +596,11 @@ class _EventParser():  # pylint:disable=too-few-public-methods
                     continue
                 if event.summary.value[0].tag == "keras":
                     self._parse_outputs(event)
+                if get_backend() == "amd":
+                    # No model is logged for AMD so need to get loss labels from state file
+                    self._add_amd_loss_labels(session_id)
                 if event.summary.value[0].tag.startswith("batch_"):
-                    data[event.step] = self._process_event(event, data.get(event.step, dict()))
+                    data[event.step] = self._process_event(event, data.get(event.step, {}))
 
         except tf_errors.DataLossError as err:
             logger.warning("The logs for Session %s are corrupted and cannot be displayed. "
@@ -604,7 +630,6 @@ class _EventParser():  # pylint:disable=too-few-public-methods
 
         config = serializer.unmarshal(struct)["config"]
         model_outputs = self._get_outputs(config)
-        split_output = len(np.unique(model_outputs[..., 1])) == 1
 
         for side_outputs, side in zip(model_outputs, ("a", "b")):
             logger.debug("side: '%s', outputs: '%s'", side, side_outputs)
@@ -615,8 +640,10 @@ class _EventParser():  # pylint:disable=too-few-public-methods
             layer_outputs = self._get_outputs(output_config)
             for output in layer_outputs:  # Drill into sub-model to get the actual output names
                 loss_name = output[0][0]
-                if not split_output:  # Rename losses to reflect the side's output
-                    loss_name = f"{loss_name.replace('_both', '')}_{side}"
+                if loss_name[-2:] not in ("_a", "_b"):  # Rename losses to reflect the side output
+                    new_name = f"{loss_name.replace('_both', '')}_{side}"
+                    logger.debug("Renaming loss output from '%s' to '%s'", loss_name, new_name)
+                    loss_name = new_name
                 if loss_name not in self._loss_labels:
                     logger.debug("Adding loss name: '%s'", loss_name)
                     self._loss_labels.append(loss_name)
@@ -647,6 +674,28 @@ class _EventParser():  # pylint:disable=too-few-public-methods
                          outputs, outputs.shape)
         return outputs
 
+    def _add_amd_loss_labels(self, session_id):
+        """ It is not possible to store the model config in the Tensorboard logs for AMD so we
+        need to obtain the loss labels from the model's state file. This is called now so we know
+        event data is being written, and therefore the most current loss label data is available
+        in the state file.
+
+        Loss names are added to :attr:`_loss_labels`
+
+        Parameters
+        ----------
+        session_id: int
+            The session id that the data is being cached for
+
+        """
+        if self._cache._loss_labels:  # pylint:disable=protected-access
+            return
+        # Import global session here to prevent circular import
+        from . import Session  # pylint:disable=import-outside-toplevel
+        loss_labels = sorted(Session.get_loss_keys(session_id=session_id))
+        self._loss_labels = loss_labels
+        logger.debug("Collated loss labels: %s", self._loss_labels)
+
     @classmethod
     def _process_event(cls, event, step):
         """ Process a single Tensorflow event.
@@ -667,8 +716,19 @@ class _EventParser():  # pylint:disable=too-few-public-methods
             The given step `dict` with the given event data added to it.
         """
         summary = event.summary.value[0]
+
         if summary.tag in ("batch_loss", "batch_total"):  # Pre tf2.3 totals were "batch_total"
             step["timestamp"] = event.wall_time
             return step
-        step.setdefault("loss", list()).append(summary.simple_value)
+
+        loss = summary.simple_value
+        if not loss:
+            # Need to convert a tensor to a float for TF2.8 logged data. This maybe due to change
+            # in logging or may be due to work around put in place in FS training function for the
+            # following bug in TF 2.8 when writing records:
+            #  https://github.com/keras-team/keras/issues/16173
+            loss = float(tf.make_ndarray(summary.tensor))
+
+        step.setdefault("loss", []).append(loss)
+
         return step

lib/gui/analysis/stats.py

@@ -17,6 +17,7 @@ from threading import Event
 import numpy as np
 
 from lib.serializer import get_serializer
+from lib.utils import get_backend
 
 from .event_reader import TensorBoardLogs
 
@@ -62,9 +63,9 @@ class GlobalSession():
     def batch_sizes(self):
         """ dict: The batch sizes for each session_id for the model. """
         if self._state is None:
-            return dict()
+            return {}
         return {int(sess_id): sess["batchsize"]
-                for sess_id, sess in self._state.get("sessions", dict()).items()}
+                for sess_id, sess in self._state.get("sessions", {}).items()}
 
     @property
     def full_summary(self):
@@ -86,7 +87,7 @@ class GlobalSession():
 
     def _load_state_file(self):
         """ Load the current state file to :attr:`_state`. """
-        state_file = os.path.join(self._model_dir, "{}_state.json".format(self._model_name))
+        state_file = os.path.join(self._model_dir, f"{self._model_name}_state.json")
         logger.debug("Loading State: '%s'", state_file)
         serializer = get_serializer("json")
         self._state = serializer.load(state_file)
@@ -125,8 +126,7 @@ class GlobalSession():
         self._model_dir = model_folder
         self._model_name = model_name
         self._load_state_file()
-        self._tb_logs = TensorBoardLogs(os.path.join(self._model_dir,
-                                                     "{}_logs".format(self._model_name)),
+        self._tb_logs = TensorBoardLogs(os.path.join(self._model_dir, f"{self._model_name}_logs"),
                                         is_training)
 
         self._summary = SessionsSummary(self)
@@ -140,7 +140,7 @@ class GlobalSession():
 
     def clear(self):
         """ Clear the currently loaded session. """
-        self._state = dict()
+        self._state = {}
         self._model_dir = None
         self._model_name = None
 
@@ -173,13 +173,13 @@ class GlobalSession():
 
         loss_dict = self._tb_logs.get_loss(session_id=session_id)
         if session_id is None:
-            retval = dict()
+            retval = {}
             for key in sorted(loss_dict):
                 for loss_key, loss in loss_dict[key].items():
                     retval.setdefault(loss_key, []).extend(loss)
             retval = {key: np.array(val, dtype="float32") for key, val in retval.items()}
         else:
-            retval = loss_dict.get(session_id, dict())
+            retval = loss_dict.get(session_id, {})
 
         if self._is_training:
             self._is_querying.clear()
@@ -239,14 +239,21 @@ class GlobalSession():
             The loss keys for the given session. If ``None`` is passed as session_id then a unique
             list of all loss keys for all sessions is returned
         """
+        if get_backend() == "amd":
+            # We can't log the graph in Tensorboard logs for AMD so need to obtain from state file
+            loss_keys = {int(sess_id): [name for name in session["loss_names"] if name != "total"]
+                         for sess_id, session in self._state["sessions"].items()}
+        else:
             loss_keys = {sess_id: list(logs.keys())
-                     for sess_id, logs in self._tb_logs.get_loss(session_id=session_id).items()}
+                         for sess_id, logs
+                         in self._tb_logs.get_loss(session_id=session_id).items()}
+
         if session_id is None:
             retval = list(set(loss_key
                               for session in loss_keys.values()
                               for loss_key in session))
         else:
-            retval = loss_keys[session_id]
+            retval = loss_keys.get(session_id)
         return retval
 
 
@@ -334,7 +341,7 @@ class SessionsSummary():  # pylint:disable=too-few-public-methods
         """
         if self._per_session_stats is None:
             logger.debug("Collating per session stats")
-            compiled = list()
+            compiled = []
             for session_id, ts_data in self._time_stats.items():
                 logger.debug("Compiling session ID: %s", session_id)
                 if self._state is None:
@@ -446,15 +453,15 @@ class SessionsSummary():  # pylint:disable=too-few-public-methods
         retval = []
         for summary in compiled_stats:
             hrs, mins, secs = self._convert_time(summary["elapsed"])
-            stats = dict()
+            stats = {}
             for key in summary:
                 if key not in ("start", "end", "elapsed", "rate"):
                     stats[key] = summary[key]
                     continue
                 stats["start"] = time.strftime("%x %X", time.localtime(summary["start"]))
                 stats["end"] = time.strftime("%x %X", time.localtime(summary["end"]))
-                stats["elapsed"] = "{}:{}:{}".format(hrs, mins, secs)
-                stats["rate"] = "{0:.1f}".format(summary["rate"])
+                stats["elapsed"] = f"{hrs}:{mins}:{secs}"
+                stats["rate"] = f"{summary['rate']:.1f}"
             retval.append(stats)
         return retval
 
@@ -474,9 +481,9 @@ class SessionsSummary():  # pylint:disable=too-few-public-methods
         """
         hrs = int(timestamp // 3600)
         if hrs < 10:
-            hrs = "{0:02d}".format(hrs)
-        mins = "{0:02d}".format((int(timestamp % 3600) // 60))
-        secs = "{0:02d}".format((int(timestamp % 3600) % 60))
+            hrs = f"{hrs:02d}"
+        mins = f"{(int(timestamp % 3600) // 60):02d}"
+        secs = f"{(int(timestamp % 3600) % 60):02d}"
         return hrs, mins, secs
 
 
@@ -529,7 +536,7 @@ class Calculations():
         self._iterations = 0
         self._limit = 0
         self._start_iteration = 0
-        self._stats = dict()
+        self._stats = {}
         self.refresh()
         logger.debug("Initialized %s", self.__class__.__name__)
 
@@ -630,7 +637,7 @@ class Calculations():
                 if self._args["flatten_outliers"]:
                     loss = self._flatten_outliers(loss)
 
-                self.stats["raw_{}".format(loss_name)] = loss
+                self.stats[f"raw_{loss_name}"] = loss
 
             self._iterations = 0 if not iterations else min(iterations)
             if self._limit > 1:
@@ -642,7 +649,7 @@ class Calculations():
             if len(iterations) > 1:
                 # Crop all losses to the same number of items
                 if self._iterations == 0:
-                    self.stats = {lossname: np.array(list(), dtype=loss.dtype)
+                    self.stats = {lossname: np.array([], dtype=loss.dtype)
                                   for lossname, loss in self.stats.items()}
                 else:
                     self.stats = {lossname: loss[:self._iterations]
@@ -722,7 +729,7 @@ class Calculations():
         logger.debug("Calculating totals rate")
         batchsizes = _SESSION.batch_sizes
         total_timestamps = _SESSION.get_timestamps(None)
-        rate = list()
+        rate = []
         for sess_id in sorted(total_timestamps.keys()):
             batchsize = batchsizes[sess_id]
             timestamps = total_timestamps[sess_id]
@@ -737,10 +744,10 @@ class Calculations():
             if selection == "raw":
                 continue
             logger.debug("Calculating: %s", selection)
-            method = getattr(self, "_calc_{}".format(selection))
+            method = getattr(self, f"_calc_{selection}")
             raw_keys = [key for key in self._stats if key.startswith("raw_")]
             for key in raw_keys:
-                selected_key = "{}_{}".format(selection, key.replace("raw_", ""))
+                selected_key = f"{selection}_{key.replace('raw_', '')}"
                 self._stats[selected_key] = method(self._stats[key])
 
     def _calc_avg(self, data):
@@ -866,7 +873,7 @@ class _ExponentialMovingAverage():  # pylint:disable=too-few-public-methods
         optimizations.
         """
         # Use :func:`np.finfo(dtype).eps` if you are worried about accuracy and want to be safe.
-        epsilon = np.finfo(self._dtype).tiny
+        epsilon = np.finfo(self._dtype).tiny  # pylint:disable=no-member
         # If this produces an OverflowError, make epsilon larger:
         retval = int(np.log(epsilon) / np.log(1 - self._alpha)) + 1
         logger.debug("row_size: %s", retval)

lib/gui/display_command.py

@@ -63,13 +63,10 @@ class PreviewExtract(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
             return
         filename = "extract_convert_preview"
         now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = os.path.join(location,
-                                "{}_{}.{}".format(filename,
-                                                  now,
-                                                  "png"))
+        filename = os.path.join(location, f"{filename}_{now}.png")
         get_images().previewoutput[0].save(filename)
         logger.debug("Saved preview to %s", filename)
-        print("Saved preview to {}".format(filename))
+        print(f"Saved preview to {filename}")
 
 
 class PreviewTrain(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
@@ -125,7 +122,7 @@ class PreviewTrain(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
         should_update = self.update_preview.get()
 
         for name in sortednames:
-            if name not in existing.keys():
+            if name not in existing:
                 self.add_child(name)
             elif should_update:
                 tab_id = existing[name]
@@ -197,19 +194,16 @@ class PreviewTrainCanvas(ttk.Frame):  # pylint: disable=too-many-ancestors
         """ Save the figure to file """
         filename = self.name
         now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = os.path.join(location,
-                                "{}_{}.{}".format(filename,
-                                                  now,
-                                                  "png"))
+        filename = os.path.join(location, f"{filename}_{now}.png")
         get_images().previewtrain[self.name][0].save(filename)
         logger.debug("Saved preview to %s", filename)
-        print("Saved preview to {}".format(filename))
+        print(f"Saved preview to {filename}")
 
 
 class GraphDisplay(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
     """ The Graph Tab of the Display section """
     def __init__(self, parent, tab_name, helptext, wait_time, command=None):
-        self._trace_vars = dict()
+        self._trace_vars = {}
         super().__init__(parent, tab_name, helptext, wait_time, command)
 
     def set_vars(self):
@@ -370,6 +364,8 @@ class GraphDisplay(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
             logger.trace("Loading graph")
             self.display_item = Session
             self._add_trace_variables()
+        elif Session.is_training and self.display_item is not None:
+            logger.trace("Graph already displayed. Nothing to do.")
         else:
             logger.trace("Clearing graph")
             self.display_item = None
@@ -384,9 +380,15 @@ class GraphDisplay(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
 
         logger.debug("Adding graph")
         existing = list(self.subnotebook_get_titles_ids().keys())
-        loss_keys = [key
-                     for key in self.display_item.get_loss_keys(Session.session_ids[-1])
-                     if key != "total"]
+
+        loss_keys = self.display_item.get_loss_keys(Session.session_ids[-1])
+        if not loss_keys:
+            # Reload if we attempt to get loss keys before data is written
+            logger.debug("Waiting for Session Data to become available to graph")
+            self.after(1000, self.display_item_process)
+            return
+
+        loss_keys = [key for key in loss_keys if key != "total"]
         display_tabs = sorted(set(key[:-1].rstrip("_") for key in loss_keys))
 
         for loss_key in display_tabs:
@@ -472,7 +474,7 @@ class GraphDisplay(DisplayOptionalPage):  # pylint: disable=too-many-ancestors
             for name, (var, trace) in self._trace_vars.items():
                 logger.debug("Clearing trace from variable: %s", name)
                 var.trace_vdelete("w", trace)
-            self._trace_vars = dict()
+            self._trace_vars = {}
 
     def close(self):
         """ Clear the plots from RAM """

lib/gui/display_page.py

@@ -59,7 +59,7 @@ class DisplayPage(ttk.Frame):  # pylint: disable=too-many-ancestors
     @staticmethod
     def set_vars():
         """ Override to return a dict of page specific variables """
-        return dict()
+        return {}
 
     def on_tab_select(self):  # pylint:disable=no-self-use
         """ Override for specific actions when the current tab is selected """
@@ -151,7 +151,7 @@ class DisplayPage(ttk.Frame):  # pylint: disable=too-many-ancestors
 
     def subnotebook_get_titles_ids(self):
         """ Return tabs ids and titles """
-        tabs = dict()
+        tabs = {}
         for tab_id in range(0, self.subnotebook.index("end")):
             tabs[self.subnotebook.tab(tab_id, "text")] = tab_id
         logger.debug(tabs)
@@ -213,11 +213,11 @@ class DisplayOptionalPage(DisplayPage):  # pylint: disable=too-many-ancestors
     def set_info_text(self):
         """ Set waiting for display text """
         if not self.vars["enabled"].get():
-            msg = "{} disabled".format(self.tabname.title())
+            msg = f"{self.tabname.title()} disabled"
         elif self.vars["enabled"].get() and not self.vars["ready"].get():
-            msg = "Waiting for {}...".format(self.tabname)
+            msg = f"Waiting for {self.tabname}..."
         else:
-            msg = "Displaying {}".format(self.tabname)
+            msg = f"Displaying {self.tabname}"
         logger.debug(msg)
         self.set_info(msg)
 
@@ -235,7 +235,7 @@ class DisplayOptionalPage(DisplayPage):  # pylint: disable=too-many-ancestors
                              command=self.save_items)
         btnsave.pack(padx=2, side=tk.RIGHT)
         Tooltip(btnsave,
-                text=_("Save {}(s) to file").format(self.tabname),
+                text=_(f"Save {self.tabname}(s) to file"),
                 wrap_length=200)
 
     def add_option_enable(self):
@@ -243,11 +243,11 @@ class DisplayOptionalPage(DisplayPage):  # pylint: disable=too-many-ancestors
         logger.debug("Adding enable option")
         chkenable = ttk.Checkbutton(self.optsframe,
                                     variable=self.vars["enabled"],
-                                    text="Enable {}".format(self.tabname),
+                                    text=f"Enable {self.tabname}",
                                     command=self.on_chkenable_change)
         chkenable.pack(side=tk.RIGHT, padx=5, anchor=tk.W)
         Tooltip(chkenable,
-                text=_("Enable or disable {} display").format(self.tabname),
+                text=_(f"Enable or disable {self.tabname} display"),
                 wrap_length=200)
 
     def save_items(self):

lib/gui/utils.py

@@ -137,6 +137,7 @@ class FileHandler():  # pylint:disable=too-few-public-methods
                      "variable: %s)", self.__class__.__name__, handle_type, file_type, title,
                      initial_folder, initial_file, command, action, variable)
         self._handletype = handle_type
+        self._dummy_master = self._set_dummy_master()
         self._defaults = self._set_defaults()
         self._kwargs = self._set_kwargs(title,
                                         initial_folder,
@@ -145,7 +146,9 @@ class FileHandler():  # pylint:disable=too-few-public-methods
                                         command,
                                         action,
                                         variable)
-        self.return_file = getattr(self, "_{}".format(self._handletype.lower()))()
+        self.return_file = getattr(self, f"_{self._handletype.lower()}")()
+        self._remove_dummy_master()
+
         logger.debug("Initialized %s", self.__class__.__name__)
 
     @property
@@ -184,10 +187,10 @@ class FileHandler():  # pylint:disable=too-few-public-methods
             if platform.system() == "Linux":
                 filetypes[key] = [item
                                   if item[0] == "All files"
-                                  else (item[0], "{} {}".format(item[1], item[1].upper()))
+                                  else (item[0], f"{item[1]} {item[1].upper()}")
                                   for item in filetypes[key]]
             if len(filetypes[key]) > 2:
-                multi = ["{} Files".format(key.title())]
+                multi = [f"{key.title()} Files"]
                 multi.append(" ".join([ftype[1]
                                        for ftype in filetypes[key] if ftype[0] != "All files"]))
                 filetypes[key].insert(0, tuple(multi))
@@ -214,6 +217,35 @@ class FileHandler():  # pylint:disable=too-few-public-methods
                                          "rotate": "save_filename",
                                          "slice": "save_filename"}))
 
+    @classmethod
+    def _set_dummy_master(cls):
+        """ Add an option to force black font on Linux file dialogs KDE issue that displays light
+        font on white background).
+
+        This is a pretty hacky solution, but tkinter does not allow direct editing of file dialogs,
+        so we create a dummy frame and add the foreground option there, so that the file dialog can
+        inherit the foreground.
+
+        Returns
+        -------
+        tkinter.Frame or ``None``
+            The dummy master frame for Linux systems, otherwise ``None``
+        """
+        if platform.system().lower() == "linux":
+            retval = tk.Frame()
+            retval.option_add("*foreground", "black")
+        else:
+            retval = None
+        return retval
+
+    def _remove_dummy_master(self):
+        """ Destroy the dummy master widget on Linux systems. """
+        if platform.system().lower() != "linux":
+            return
+        self._dummy_master.destroy()
+        del self._dummy_master
+        self._dummy_master = None
+
     def _set_defaults(self):
         """ Set the default file type for the file dialog. Generally the first found file type
         will be used, but this is overridden if it is not appropriate.
@@ -264,7 +296,9 @@ class FileHandler():  # pylint:disable=too-few-public-methods
         logger.debug("Setting Kwargs: (title: %s, initial_folder: %s, initial_file: '%s', "
                      "file_type: '%s', command: '%s': action: '%s', variable: '%s')",
                      title, initial_folder, initial_file, file_type, command, action, variable)
-        kwargs = dict()
+
+        kwargs = dict(master=self._dummy_master)
+
         if self._handletype.lower() == "context":
             self._set_context_handletype(command, action, variable)
 
@@ -361,10 +395,10 @@ class Images():
         self._pathpreview = os.path.join(PATHCACHE, "preview")
         self._pathoutput = None
         self._previewoutput = None
-        self._previewtrain = dict()
+        self._previewtrain = {}
         self._previewcache = dict(modified=None,  # cache for extract and convert
                                   images=None,
-                                  filenames=list(),
+                                  filenames=[],
                                   placeholder=None)
         self._errcount = 0
         self._icons = self._load_icons()
@@ -420,7 +454,7 @@ class Images():
         """
         size = get_config().user_config_dict.get("icon_size", 16)
         size = int(round(size * get_config().scaling_factor))
-        icons = dict()
+        icons = {}
         pathicons = os.path.join(PATHCACHE, "icons")
         for fname in os.listdir(pathicons):
             name, ext = os.path.splitext(fname)
@@ -470,10 +504,10 @@ class Images():
         logger.debug("Clearing image cache")
         self._pathoutput = None
         self._previewoutput = None
-        self._previewtrain = dict()
+        self._previewtrain = {}
         self._previewcache = dict(modified=None,  # cache for extract and convert
                                   images=None,
-                                  filenames=list(),
+                                  filenames=[],
                                   placeholder=None)
 
     @staticmethod
@@ -600,10 +634,10 @@ class Images():
         logger.debug("num_images: %s", num_images)
         if num_images == 0:
             return False
-        samples = list()
+        samples = []
         start_idx = len(image_files) - num_images if len(image_files) > num_images else 0
         show_files = sorted(image_files, key=os.path.getctime)[start_idx:]
-        dropped_files = list()
+        dropped_files = []
         for fname in show_files:
             try:
                 img = Image.open(fname)
@@ -732,7 +766,7 @@ class Images():
         modified = None
         if not image_files:
             logger.debug("No preview to display")
-            self._previewtrain = dict()
+            self._previewtrain = {}
             return
         for img in image_files:
             modified = os.path.getmtime(img) if modified is None else modified
@@ -755,7 +789,7 @@ class Images():
                     self._errcount += 1
                 else:
                     logger.error("Error reading the preview file for '%s'", img)
-                    print("Error reading the preview file for {}".format(name))
+                    print(f"Error reading the preview file for {name}")
                     self._previewtrain[name] = None
 
     def _get_current_size(self, name):
@@ -1126,7 +1160,7 @@ class Config():
             Additional text to be appended to the GUI title bar. Default: ``None``
         """
         title = "Faceswap.py"
-        title += " - {}".format(text) if text is not None and text else ""
+        title += f" - {text}" if text is not None and text else ""
         self.root.title(title)
 
     def set_geometry(self, width, height, fullscreen=False):
@@ -1154,8 +1188,7 @@ class Config():
         elif fullscreen:
             self.root.attributes('-zoomed', True)
         else:
-            self.root.geometry("{}x{}+80+80".format(str(initial_dimensions[0]),
-                                                    str(initial_dimensions[1])))
+            self.root.geometry(f"{str(initial_dimensions[0])}x{str(initial_dimensions[1])}+80+80")
         logger.debug("Geometry: %sx%s", *initial_dimensions)
 
 
@@ -1260,7 +1293,7 @@ class PreviewTrigger():
          """
         trigger = self._trigger_files[trigger_type]
         if not os.path.isfile(trigger):
-            with open(trigger, "w"):
+            with open(trigger, "w", encoding="utf8"):
                 pass
             logger.debug("Set preview trigger: %s", trigger)
 

lib/keypress.py

@@ -17,6 +17,7 @@ GNU General Public License for more details.
 """
 
 import os
+import sys
 
 # Windows
 if os.name == "nt":
@@ -24,7 +25,6 @@ if os.name == "nt":
 
 # Posix (Linux, OS X)
 else:
-    import sys
     import termios
     import atexit
     from select import select
@@ -34,7 +34,7 @@ class KBHit:
     """ Creates a KBHit object that you can call to do various keyboard things. """
     def __init__(self, is_gui=False):
         self.is_gui = is_gui
-        if os.name == "nt" or self.is_gui:
+        if os.name == "nt" or self.is_gui or not sys.stdout.isatty():
             pass
         else:
             # Save the terminal settings
@@ -51,7 +51,7 @@ class KBHit:
 
     def set_normal_term(self):
         """ Resets to normal terminal.  On Windows this is a no-op. """
-        if os.name == "nt" or self.is_gui:
+        if os.name == "nt" or self.is_gui or not sys.stdout.isatty():
             pass
         else:
             termios.tcsetattr(self.file_desc, termios.TCSAFLUSH, self.old_term)
@@ -59,7 +59,7 @@ class KBHit:
     def getch(self):
         """ Returns a keyboard character after kbhit() has been called.
             Should not be called in the same program as getarrow(). """
-        if self.is_gui and os.name != "nt":
+        if self.is_gui and os.name != "nt" or not sys.stdout.isatty():
             return None
         if os.name == "nt":
             return msvcrt.getch().decode("utf-8")
@@ -73,7 +73,7 @@ class KBHit:
         3 : left
         Should not be called in the same program as getch(). """
 
-        if self.is_gui:
+        if self.is_gui or not sys.stdout.isatty():
             return None
         if os.name == "nt":
             msvcrt.getch()  # skip 0xE0
@@ -87,7 +87,7 @@ class KBHit:
 
     def kbhit(self):
         """ Returns True if keyboard character was hit, False otherwise. """
-        if self.is_gui and os.name != "nt":
+        if self.is_gui and os.name != "nt" or not sys.stdout.isatty():
             return None
         if os.name == "nt":
             return msvcrt.kbhit()

lib/model/initializers.py

@@ -7,16 +7,18 @@ import inspect
 
 import numpy as np
 import tensorflow as tf
-from keras import backend as K
-from keras import initializers
-
-try:
-    from keras.utils import get_custom_objects
-except ImportError:
-    from tensorflow.keras.utils import get_custom_objects
 
 from lib.utils import get_backend
 
+if get_backend() == "amd":
+    from keras.utils import get_custom_objects  # pylint:disable=no-name-in-module
+    from keras import backend as K
+    from keras import initializers
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.utils import get_custom_objects  # noqa pylint:disable=no-name-in-module,import-error
+    from tensorflow.keras import initializers, backend as K  # noqa pylint:disable=no-name-in-module,import-error
+
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
 
@@ -68,7 +70,7 @@ def compute_fans(shape, data_format='channels_last'):
     return fan_in, fan_out
 
 
-class ICNR(initializers.Initializer):  # pylint: disable=invalid-name
+class ICNR(initializers.Initializer):  # pylint: disable=invalid-name,no-member
     """ ICNR initializer for checkerboard artifact free sub pixel convolution
 
     Parameters
@@ -171,11 +173,11 @@ class ICNR(initializers.Initializer):  # pylint: disable=invalid-name
         config = {"scale": self.scale,
                   "initializer": self.initializer
                   }
-        base_config = super(ICNR, self).get_config()
+        base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
-class ConvolutionAware(initializers.Initializer):
+class ConvolutionAware(initializers.Initializer):  # pylint: disable=no-member
     """
     Initializer that generates orthogonal convolution filters in the Fourier space. If this
     initializer is passed a shape that is not 3D or 4D, orthogonal initialization will be used.
@@ -208,8 +210,8 @@ class ConvolutionAware(initializers.Initializer):
     def __init__(self, eps_std=0.05, seed=None, initialized=False):
         self.eps_std = eps_std
         self.seed = seed
-        self.orthogonal = initializers.Orthogonal()
-        self.he_uniform = initializers.he_uniform()
+        self.orthogonal = initializers.Orthogonal()  # pylint:disable=no-member
+        self.he_uniform = initializers.he_uniform()  # pylint:disable=no-member
         self.initialized = initialized
 
     def __call__(self, shape, dtype=None):

lib/model/layers.py

@@ -7,22 +7,21 @@ import sys
 import inspect
 
 import tensorflow as tf
-import keras.backend as K
-
-from keras.layers import InputSpec, Layer
-try:
-    from keras.utils import get_custom_objects
-except ImportError:
-    from tensorflow.keras.utils import get_custom_objects
 
 from lib.utils import get_backend
 
 if get_backend() == "amd":
     from lib.plaidml_utils import pad
-    from keras.utils import conv_utils  # pylint:disable=ungrouped-imports
+    from keras.utils import get_custom_objects, conv_utils  # pylint:disable=no-name-in-module
+    import keras.backend as K
+    from keras.layers import InputSpec, Layer
 else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.utils import get_custom_objects  # noqa pylint:disable=no-name-in-module,import-error
+    from tensorflow.keras import backend as K  # pylint:disable=import-error
+    from tensorflow.keras.layers import InputSpec, Layer  # noqa pylint:disable=no-name-in-module,import-error
     from tensorflow import pad
-    from tensorflow.python.keras.utils import conv_utils
+    from tensorflow.python.keras.utils import conv_utils  # pylint:disable=no-name-in-module
 
 
 class PixelShuffler(Layer):
@@ -67,20 +66,22 @@ class PixelShuffler(Layer):
     def __init__(self, size=(2, 2), data_format=None, **kwargs):
         super().__init__(**kwargs)
         if get_backend() == "amd":
-            self.data_format = K.normalize_data_format(data_format)
+            self.data_format = K.normalize_data_format(data_format)  # pylint:disable=no-member
         else:
             self.data_format = conv_utils.normalize_data_format(data_format)
         self.size = conv_utils.normalize_tuple(size, 2, 'size')
 
-    def call(self, inputs, **kwargs):  # pylint:disable=unused-argument
+    def call(self, inputs, *args, **kwargs):
         """This is where the layer's logic lives.
 
         Parameters
         ----------
         inputs: tensor
             Input tensor, or list/tuple of input tensors
+        args: tuple
+            Additional standard keras Layer arguments
         kwargs: dict
-            Additional keyword arguments. Unused
+            Additional standard keras Layer keyword arguments
 
         Returns
         -------
@@ -189,7 +190,7 @@ class PixelShuffler(Layer):
         """
         config = {'size': self.size,
                   'data_format': self.data_format}
-        base_config = super(PixelShuffler, self).get_config()
+        base_config = super().get_config()
 
         return dict(list(base_config.items()) + list(config.items()))
 
@@ -211,15 +212,17 @@ class KResizeImages(Layer):
         self.size = size
         self.interpolation = interpolation
 
-    def call(self, inputs, **kwargs):  # pylint:disable=unused-argument
+    def call(self, inputs, *args, **kwargs):
         """ Call the upsample layer
 
         Parameters
         ----------
         inputs: tensor
             Input tensor, or list/tuple of input tensors
+        args: tuple
+            Additional standard keras Layer arguments
         kwargs: dict
-            Additional keyword arguments. Unused
+            Additional standard keras Layer keyword arguments
 
         Returns
         -------
@@ -319,11 +322,11 @@ class SubPixelUpscaling(Layer):
     """
 
     def __init__(self, scale_factor=2, data_format=None, **kwargs):
-        super(SubPixelUpscaling, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
         self.scale_factor = scale_factor
         if get_backend() == "amd":
-            self.data_format = K.normalize_data_format(data_format)
+            self.data_format = K.normalize_data_format(data_format)  # pylint:disable=no-member
         else:
             self.data_format = conv_utils.normalize_data_format(data_format)
 
@@ -340,15 +343,17 @@ class SubPixelUpscaling(Layer):
         """
         pass  # pylint: disable=unnecessary-pass
 
-    def call(self, inputs, **kwargs):  # pylint:disable=unused-argument
+    def call(self, inputs, *args, **kwargs):
         """This is where the layer's logic lives.
 
         Parameters
         ----------
         inputs: tensor
             Input tensor, or list/tuple of input tensors
+        args: tuple
+            Additional standard keras Layer arguments
         kwargs: dict
-            Additional keyword arguments. Unused
+            Additional standard keras Layer keyword arguments
 
         Returns
         -------
@@ -463,7 +468,7 @@ class SubPixelUpscaling(Layer):
         """
         config = {"scale_factor": self.scale_factor,
                   "data_format": self.data_format}
-        base_config = super(SubPixelUpscaling, self).get_config()
+        base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
@@ -595,7 +600,7 @@ class ReflectionPadding2D(Layer):
         """
         config = {'stride': self.stride,
                   'kernel_size': self.kernel_size}
-        base_config = super(ReflectionPadding2D, self).get_config()
+        base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
@@ -605,9 +610,9 @@ class _GlobalPooling2D(Layer):
     From keras as access to pooling is trickier in tensorflow.keras
     """
     def __init__(self, data_format=None, **kwargs):
-        super(_GlobalPooling2D, self).__init__(**kwargs)
+        super().__init__(**kwargs)
         if get_backend() == "amd":
-            self.data_format = K.normalize_data_format(data_format)
+            self.data_format = K.normalize_data_format(data_format)  # pylint:disable=no-member
         else:
             self.data_format = conv_utils.normalize_data_format(data_format)
         self.input_spec = InputSpec(ndim=4)
@@ -624,37 +629,41 @@ class _GlobalPooling2D(Layer):
             return (input_shape[0], input_shape[3])
         return (input_shape[0], input_shape[1])
 
-    def call(self, inputs, **kwargs):
+    def call(self, inputs, *args, **kwargs):
         """ Override to call the layer.
 
         Parameters
         ----------
         inputs: Tensor
             The input to the layer
+        args: tuple
+            Additional standard keras Layer arguments
         kwargs: dict
-            Additional keyword arguments
+            Additional standard keras Layer keyword arguments
         """
         raise NotImplementedError
 
     def get_config(self):
         """ Set the Keras config """
         config = {'data_format': self.data_format}
-        base_config = super(_GlobalPooling2D, self).get_config()
+        base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
 class GlobalMinPooling2D(_GlobalPooling2D):
     """Global minimum pooling operation for spatial data. """
 
-    def call(self, inputs, **kwargs):
+    def call(self, inputs, *args, **kwargs):
         """This is where the layer's logic lives.
 
         Parameters
         ----------
         inputs: tensor
             Input tensor, or list/tuple of input tensors
+        args: tuple
+            Additional standard keras Layer arguments
         kwargs: dict
-            Additional keyword arguments
+            Additional standard keras Layer keyword arguments
 
         Returns
         -------
@@ -671,15 +680,17 @@ class GlobalMinPooling2D(_GlobalPooling2D):
 class GlobalStdDevPooling2D(_GlobalPooling2D):
     """Global standard deviation pooling operation for spatial data. """
 
-    def call(self, inputs, **kwargs):
+    def call(self, inputs, *args, **kwargs):
         """This is where the layer's logic lives.
 
         Parameters
         ----------
         inputs: tensor
             Input tensor, or list/tuple of input tensors
+        args: tuple
+            Additional standard keras Layer arguments
         kwargs: dict
-            Additional keyword arguments
+            Additional standard keras Layer keyword arguments
 
         Returns
         -------
@@ -705,7 +716,7 @@ class L2_normalize(Layer):  # pylint:disable=invalid-name
     """
     def __init__(self, axis, **kwargs):
         self.axis = axis
-        super(L2_normalize, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
     def call(self, inputs):  # pylint:disable=arguments-differ
         """This is where the layer's logic lives.
@@ -739,7 +750,7 @@ class L2_normalize(Layer):  # pylint:disable=invalid-name
         dict
             A python dictionary containing the layer configuration
         """
-        config = super(L2_normalize, self).get_config()
+        config = super().get_config()
         config["axis"] = self.axis
         return config
 

lib/model/losses_plaid.py

@@ -199,6 +199,64 @@ class DSSIMObjective():
         return patches
 
 
+class MSSSIMLoss():  # pylint:disable=too-few-public-methods
+    """ Multiscale Structural Similarity Loss Function
+
+    Parameters
+    ----------
+    k_1: float, optional
+        Parameter of the SSIM. Default: `0.01`
+    k_2: float, optional
+        Parameter of the SSIM. Default: `0.03`
+    filter_size: int, optional
+        size of gaussian filter Default: `11`
+    filter_sigma: float, optional
+        Width of gaussian filter Default: `1.5`
+    max_value: float, optional
+        Max value of the output. Default: `1.0`
+    power_factors: tuple, optional
+        Iterable of weights for each of the scales. The number of scales used is the length of the
+        list. Index 0 is the unscaled resolution's weight and each increasing scale corresponds to
+        the image being downsampled by 2. Defaults to the values obtained in the original paper.
+        Default: (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
+
+    Notes
+    ------
+    You should add a regularization term like a l2 loss in addition to this one.
+    """
+    def __init__(self,
+                 k_1=0.01,
+                 k_2=0.03,
+                 filter_size=4,
+                 filter_sigma=1.5,
+                 max_value=1.0,
+                 power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)):
+        self.filter_size = filter_size
+        self.filter_sigma = filter_sigma
+        self.k_1 = k_1
+        self.k_2 = k_2
+        self.max_value = max_value
+        self.power_factors = power_factors
+
+    def __call__(self, y_true, y_pred):
+        """ Call the MS-SSIM Loss Function.
+
+        Parameters
+        ----------
+        y_true: tensor or variable
+            The ground truth value
+        y_pred: tensor or variable
+            The predicted value
+
+        Returns
+        -------
+        tensor
+            The MS-SSIM Loss value
+        """
+        raise FaceswapError("MS-SSIM Loss is not currently compatible with PlaidML. Please select "
+                            "a different Loss method.")
+
+
 class GeneralizedLoss():  # pylint:disable=too-few-public-methods
     """  Generalized function used to return a large variety of mathematical loss functions.
 

lib/model/losses_tf.py

@@ -7,17 +7,18 @@ import logging
 
 import numpy as np
 import tensorflow as tf
-from tensorflow.python.keras.engine import compile_utils
 
-from keras import backend as K
+# Ignore linting errors from Tensorflow's thoroughly broken import system
+from tensorflow.python.keras.engine import compile_utils  # noqa pylint:disable=no-name-in-module,import-error
+from tensorflow.keras import backend as K  # pylint:disable=import-error
 
-logger = logging.getLogger(__name__)  # pylint:disable=invalid-name
+logger = logging.getLogger(__name__)
 
 
-class DSSIMObjective(tf.keras.losses.Loss):
+class DSSIMObjective(tf.keras.losses.Loss):  # pylint:disable=too-few-public-methods
     """ DSSIM Loss Function
 
-    Difference of Structural Similarity (DSSIM loss function). Clipped between 0 and 0.5
+    Difference of Structural Similarity (DSSIM loss function).
 
     Parameters
     ----------
@@ -25,66 +26,24 @@ class DSSIMObjective(tf.keras.losses.Loss):
         Parameter of the SSIM. Default: `0.01`
     k_2: float, optional
         Parameter of the SSIM. Default: `0.03`
-    kernel_size: int, optional
-        Size of the sliding window Default: `3`
+    filter_size: int, optional
+        size of gaussian filter Default: `11`
+    filter_sigma: float, optional
+        Width of gaussian filter Default: `1.5`
     max_value: float, optional
         Max value of the output. Default: `1.0`
 
     Notes
     ------
     You should add a regularization term like a l2 loss in addition to this one.
-
-    References
-    ----------
-    https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/dssim.py
-
-    MIT License
-
-    Copyright (c) 2017 Fariz Rahman
-
-    Permission is hereby granted, free of charge, to any person obtaining a copy
-    of this software and associated documentation files (the "Software"), to deal
-    in the Software without restriction, including without limitation the rights
-    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-    copies of the Software, and to permit persons to whom the Software is
-    furnished to do so, subject to the following conditions:
-
-    The above copyright notice and this permission notice shall be included in all
-    copies or substantial portions of the Software.
-
-    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-    SOFTWARE.
     """
-    def __init__(self, k_1=0.01, k_2=0.03, kernel_size=3, max_value=1.0):
+    def __init__(self, k_1=0.01, k_2=0.03, filter_size=11, filter_sigma=1.5, max_value=1.0):
         super().__init__(name="DSSIMObjective")
-        self.kernel_size = kernel_size
+        self.filter_size = filter_size
+        self.filter_sigma = filter_sigma
         self.k_1 = k_1
         self.k_2 = k_2
         self.max_value = max_value
-        self.c_1 = (self.k_1 * self.max_value) ** 2
-        self.c_2 = (self.k_2 * self.max_value) ** 2
-        self.dim_ordering = K.image_data_format()
-
-    @staticmethod
-    def _int_shape(input_tensor):
-        """ Returns the shape of tensor or variable as a tuple of int or None entries.
-
-        Parameters
-        ----------
-        input_tensor: tensor or variable
-            The input to return the shape for
-
-        Returns
-        -------
-        tuple
-            A tuple of integers (or None entries)
-        """
-        return K.int_shape(input_tensor)
 
     def call(self, y_true, y_pred):
         """ Call the DSSIM Loss Function.
@@ -100,104 +59,113 @@ class DSSIMObjective(tf.keras.losses.Loss):
         -------
         tensor
             The DSSIM Loss value
+        """
+        ssim = tf.image.ssim(y_true,
+                             y_pred,
+                             self.max_value,
+                             filter_size=self.filter_size,
+                             filter_sigma=self.filter_sigma,
+                             k1=self.k_1,
+                             k2=self.k_2)
+        dssim_loss = (1. - ssim) / 2.0
+        return dssim_loss
+
+
+class MSSSIMLoss(tf.keras.losses.Loss):  # pylint:disable=too-few-public-methods
+    """ Multiscale Structural Similarity Loss Function
+
+    Parameters
+    ----------
+    k_1: float, optional
+        Parameter of the SSIM. Default: `0.01`
+    k_2: float, optional
+        Parameter of the SSIM. Default: `0.03`
+    filter_size: int, optional
+        size of gaussian filter Default: `11`
+    filter_sigma: float, optional
+        Width of gaussian filter Default: `1.5`
+    max_value: float, optional
+        Max value of the output. Default: `1.0`
+    power_factors: tuple, optional
+        Iterable of weights for each of the scales. The number of scales used is the length of the
+        list. Index 0 is the unscaled resolution's weight and each increasing scale corresponds to
+        the image being downsampled by 2. Defaults to the values obtained in the original paper.
+        Default: (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
 
     Notes
-        -----
-        There are additional parameters for this function. some of the 'modes' for edge behavior
-        do not yet have a gradient definition in the Theano tree and cannot be used for learning
-        """
-
-        kernel = [self.kernel_size, self.kernel_size]
-        y_true = K.reshape(y_true, [-1] + list(self._int_shape(y_pred)[1:]))
-        y_pred = K.reshape(y_pred, [-1] + list(self._int_shape(y_pred)[1:]))
-        patches_pred = self.extract_image_patches(y_pred,
-                                                  kernel,
-                                                  kernel,
-                                                  'valid',
-                                                  self.dim_ordering)
-        patches_true = self.extract_image_patches(y_true,
-                                                  kernel,
-                                                  kernel,
-                                                  'valid',
-                                                  self.dim_ordering)
-
-        # Get mean
-        u_true = K.mean(patches_true, axis=-1)
-        u_pred = K.mean(patches_pred, axis=-1)
-        # Get variance
-        var_true = K.var(patches_true, axis=-1)
-        var_pred = K.var(patches_pred, axis=-1)
-        # Get standard deviation
-        covar_true_pred = K.mean(
-            patches_true * patches_pred, axis=-1) - u_true * u_pred
-
-        ssim = (2 * u_true * u_pred + self.c_1) * (
-            2 * covar_true_pred + self.c_2)
-        denom = (K.square(u_true) + K.square(u_pred) + self.c_1) * (
-            var_pred + var_true + self.c_2)
-        ssim /= denom  # no need for clipping, c_1 + c_2 make the denorm non-zero
-        return (1.0 - ssim) / 2.0
-
-    @staticmethod
-    def _preprocess_padding(padding):
-        """Convert keras padding to tensorflow padding.
-
-        Parameters
-        ----------
-        padding: string,
-            `"same"` or `"valid"`.
-
-        Returns
-        -------
-        str
-            `"SAME"` or `"VALID"`.
-
-        Raises
     ------
-        ValueError
-            If `padding` is invalid.
+    You should add a regularization term like a l2 loss in addition to this one.
     """
-        if padding == 'same':
-            padding = 'SAME'
-        elif padding == 'valid':
-            padding = 'VALID'
-        else:
-            raise ValueError('Invalid padding:', padding)
-        return padding
+    def __init__(self,
+                 k_1=0.01,
+                 k_2=0.03,
+                 filter_size=4,
+                 filter_sigma=1.5,
+                 max_value=1.0,
+                 power_factors=(0.0448, 0.2856, 0.3001, 0.2363, 0.1333)):
+        super().__init__(name="SSIM_Multiscale_Loss")
+        self.filter_size = filter_size
+        self.filter_sigma = filter_sigma
+        self.k_1 = k_1
+        self.k_2 = k_2
+        self.max_value = max_value
+        self.power_factors = power_factors
 
-    def extract_image_patches(self, input_tensor, k_sizes, s_sizes,
-                              padding='same', data_format='channels_last'):
-        """ Extract the patches from an image.
+    def call(self, y_true, y_pred):
+        """ Call the MS-SSIM Loss Function.
 
         Parameters
         ----------
-        input_tensor: tensor
-            The input image
-        k_sizes: tuple
-            2-d tuple with the kernel size
-        s_sizes: tuple
-            2-d tuple with the strides size
-        padding: str, optional
-            `"same"` or `"valid"`. Default: `"same"`
-        data_format: str, optional.
-            `"channels_last"` or `"channels_first"`. Default: `"channels_last"`
+        y_true: tensor or variable
+            The ground truth value
+        y_pred: tensor or variable
+            The predicted value
 
         Returns
         -------
-        The (k_w, k_h) patches extracted
-            Tensorflow ==> (batch_size, w, h, k_w, k_h, c)
-            Theano ==> (batch_size, w, h, c, k_w, k_h)
+        tensor
+            The MS-SSIM Loss value
         """
-        kernel = [1, k_sizes[0], k_sizes[1], 1]
-        strides = [1, s_sizes[0], s_sizes[1], 1]
-        padding = self._preprocess_padding(padding)
-        if data_format == 'channels_first':
-            input_tensor = K.permute_dimensions(input_tensor, (0, 2, 3, 1))
-        patches = tf.image.extract_patches(input_tensor, kernel, strides, [1, 1, 1, 1], padding)
-        return patches
+        im_size = K.int_shape(y_true)[1]
+        # filter size cannot be larger than the smallest scale
+        smallest_scale = self._get_smallest_size(im_size, len(self.power_factors) - 1)
+        filter_size = min(self.filter_size, smallest_scale)
+
+        ms_ssim = tf.image.ssim_multiscale(y_true,
+                                           y_pred,
+                                           self.max_value,
+                                           power_factors=self.power_factors,
+                                           filter_size=filter_size,
+                                           filter_sigma=self.filter_sigma,
+                                           k1=self.k_1,
+                                           k2=self.k_2)
+        ms_ssim_loss = 1. - ms_ssim
+        return ms_ssim_loss
+
+    def _get_smallest_size(self, size, idx):
+        """ Recursive function to obtain the smallest size that the image will be scaled to.
+
+        Parameters
+        ----------
+        size: int
+            The current scaled size to iterate through
+        idx: int
+            The current iteration to be performed. When iteration hits zero the value will
+            be returned
+
+        Returns
+        -------
+        int
+            The smallest size the image will be scaled to based on the original image size and
+            the amount of scaling factors that will occur
+        """
+        logger.debug("scale id: %s, size: %s", idx, size)
+        if idx > 0:
+            size = self._get_smallest_size(size // 2, idx - 1)
+        return size
 
 
-class GeneralizedLoss(tf.keras.losses.Loss):
+class GeneralizedLoss(tf.keras.losses.Loss):  # pylint:disable=too-few-public-methods
     """  Generalized function used to return a large variety of mathematical loss functions.
 
     The primary benefit is a smooth, differentiable version of L1 loss.
@@ -247,10 +215,11 @@ class GeneralizedLoss(tf.keras.losses.Loss):
         return loss
 
 
-class LInfNorm(tf.keras.losses.Loss):
+class LInfNorm(tf.keras.losses.Loss):  # pylint:disable=too-few-public-methods
     """ Calculate the L-inf norm as a loss function. """
 
-    def call(self, y_true, y_pred):
+    @classmethod
+    def call(cls, y_true, y_pred):
         """ Call the L-inf norm loss function.
 
         Parameters
@@ -271,7 +240,7 @@ class LInfNorm(tf.keras.losses.Loss):
         return loss
 
 
-class GradientLoss(tf.keras.losses.Loss):
+class GradientLoss(tf.keras.losses.Loss):  # pylint:disable=too-few-public-methods
     """ Gradient Loss Function.
 
     Calculates the first and second order gradient difference between pixels of an image in the x
@@ -392,7 +361,7 @@ class GradientLoss(tf.keras.losses.Loss):
         return (xy_out1 - xy_out2) * 0.25
 
 
-class GMSDLoss(tf.keras.losses.Loss):
+class GMSDLoss(tf.keras.losses.Loss):  # pylint:disable=too-few-public-methods
     """ Gradient Magnitude Similarity Deviation Loss.
 
     Improved image quality metric over MS-SSIM with easier calculations
@@ -486,7 +455,9 @@ class GMSDLoss(tf.keras.losses.Loss):
         # Use depth-wise convolution to calculate edge maps per channel.
         # Output tensor has shape [batch_size, h, w, d * num_kernels].
         pad_sizes = [[0, 0], [2, 2], [2, 2], [0, 0]]
-        padded = tf.pad(image, pad_sizes, mode='REFLECT')
+        padded = tf.pad(image,  # pylint:disable=unexpected-keyword-arg,no-value-for-parameter
+                        pad_sizes,
+                        mode='REFLECT')
         output = K.depthwise_conv2d(padded, kernels)
 
         if not magnitude:  # direction of edges

lib/model/nn_blocks.py

@@ -3,20 +3,31 @@
 
 import logging
 
-from keras.layers import (Activation, Add, BatchNormalization, Concatenate, Conv2D as KConv2D,
-                          Conv2DTranspose, DepthwiseConv2D as KDepthwiseConv2d, LeakyReLU, PReLU,
-                          SeparableConv2D, UpSampling2D)
-from keras.initializers import he_uniform, VarianceScaling
+from lib.utils import get_backend
+
 
 from .initializers import ICNR, ConvolutionAware
 from .layers import PixelShuffler, ReflectionPadding2D, Swish, KResizeImages
 from .normalization import InstanceNormalization
 
+if get_backend() == "amd":
+    from keras.layers import (
+        Activation, Add, BatchNormalization, Concatenate, Conv2D as KConv2D, Conv2DTranspose,
+        DepthwiseConv2D as KDepthwiseConv2d, LeakyReLU, PReLU, SeparableConv2D, UpSampling2D)
+    from keras.initializers import he_uniform, VarianceScaling  # pylint:disable=no-name-in-module
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import (  # noqa pylint:disable=no-name-in-module,import-error
+        Activation, Add, BatchNormalization, Concatenate, Conv2D as KConv2D, Conv2DTranspose,
+        DepthwiseConv2D as KDepthwiseConv2d, LeakyReLU, PReLU, SeparableConv2D, UpSampling2D)
+    from tensorflow.keras.initializers import he_uniform, VarianceScaling  # noqa pylint:disable=no-name-in-module,import-error
+
+
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
 
-_CONFIG = dict()
-_NAMES = dict()
+_CONFIG = {}
+_NAMES = {}
 
 
 def set_config(configuration):
@@ -52,9 +63,9 @@ def _get_name(name):
     str
         The unique name for this layer
     """
-    global _NAMES  # pylint:disable=global-statement
+    global _NAMES  # pylint:disable=global-statement,global-variable-not-assigned
     _NAMES[name] = _NAMES.setdefault(name, -1) + 1
-    name = "{}_{}".format(name, _NAMES[name])
+    name = f"{name}_{_NAMES[name]}"
     logger.debug("Generating block name: %s", name)
     return name
 
@@ -112,7 +123,7 @@ class Conv2D(KConv2D):  # pylint:disable=too-few-public-methods
     def __init__(self, *args, padding="same", check_icnr_init=False, **kwargs):
         if kwargs.get("name", None) is None:
             filters = kwargs["filters"] if "filters" in kwargs else args[0]
-            kwargs["name"] = _get_name("conv2d_{}".format(filters))
+            kwargs["name"] = _get_name(f"conv2d_{filters}")
         initializer = _get_default_initializer(kwargs.pop("kernel_initializer", None))
         if check_icnr_init and _CONFIG["icnr_init"]:
             initializer = ICNR(initializer=initializer)
@@ -179,7 +190,7 @@ class Conv2DOutput():  # pylint:disable=too-few-public-methods
     """
     def __init__(self, filters, kernel_size, activation="sigmoid", padding="same", **kwargs):
         self._name = kwargs.pop("name") if "name" in kwargs else _get_name(
-            "conv_output_{}".format(filters))
+            f"conv_output_{filters}")
         self._filters = filters
         self._kernel_size = kernel_size
         self._activation = activation
@@ -202,7 +213,7 @@ class Conv2DOutput():  # pylint:disable=too-few-public-methods
         var_x = Conv2D(self._filters,
                        self._kernel_size,
                        padding=self._padding,
-                       name="{}_conv2d".format(self._name),
+                       name=f"{self._name}_conv2d",
                        **self._kwargs)(inputs)
         var_x = Activation(self._activation, dtype="float32", name=self._name)(var_x)
         return var_x
@@ -256,8 +267,7 @@ class Conv2DBlock():  # pylint:disable=too-few-public-methods
                  activation="leakyrelu",
                  use_depthwise=False,
                  **kwargs):
-        self._name = kwargs.pop("name") if "name" in kwargs else _get_name(
-            "conv_{}".format(filters))
+        self._name = kwargs.pop("name") if "name" in kwargs else _get_name(f"conv_{filters}")
 
         logger.debug("name: %s, filters: %s, kernel_size: %s, strides: %s, padding: %s, "
                      "normalization: %s, activation: %s, use_depthwise: %s, kwargs: %s)",
@@ -299,26 +309,26 @@ class Conv2DBlock():  # pylint:disable=too-few-public-methods
         if self._use_reflect_padding:
             inputs = ReflectionPadding2D(stride=self._strides,
                                          kernel_size=self._args[-1],
-                                         name="{}_reflectionpadding2d".format(self._name))(inputs)
+                                         name=f"{self._name}_reflectionpadding2d")(inputs)
         conv = DepthwiseConv2D if self._use_depthwise else Conv2D
         var_x = conv(*self._args,
                      strides=self._strides,
                      padding=self._padding,
-                     name="{}_{}conv2d".format(self._name, "dw" if self._use_depthwise else ""),
+                     name=f"{self._name}_{'dw' if self._use_depthwise else ''}conv2d",
                      **self._kwargs)(inputs)
         # normalization
         if self._normalization == "instance":
-            var_x = InstanceNormalization(name="{}_instancenorm".format(self._name))(var_x)
+            var_x = InstanceNormalization(name=f"{self._name}_instancenorm")(var_x)
         if self._normalization == "batch":
-            var_x = BatchNormalization(axis=3, name="{}_batchnorm".format(self._name))(var_x)
+            var_x = BatchNormalization(axis=3, name=f"{self._name}_batchnorm")(var_x)
 
         # activation
         if self._activation == "leakyrelu":
-            var_x = LeakyReLU(0.1, name="{}_leakyrelu".format(self._name))(var_x)
+            var_x = LeakyReLU(0.1, name=f"{self._name}_leakyrelu")(var_x)
         if self._activation == "swish":
-            var_x = Swish(name="{}_swish".format(self._name))(var_x)
+            var_x = Swish(name=f"{self._name}_swish")(var_x)
         if self._activation == "prelu":
-            var_x = PReLU(name="{}_prelu".format(self._name))(var_x)
+            var_x = PReLU(name=f"{self._name}_prelu")(var_x)
 
         return var_x
 
@@ -344,7 +354,7 @@ class SeparableConv2DBlock():  # pylint:disable=too-few-public-methods
         Convolutional 2D layer
     """
     def __init__(self, filters, kernel_size=5, strides=2, **kwargs):
-        self._name = _get_name("separableconv2d_{}".format(filters))
+        self._name = _get_name(f"separableconv2d_{filters}")
         logger.debug("name: %s, filters: %s, kernel_size: %s, strides: %s, kwargs: %s)",
                      self._name, filters, kernel_size, strides, kwargs)
 
@@ -373,9 +383,9 @@ class SeparableConv2DBlock():  # pylint:disable=too-few-public-methods
                                 kernel_size=self._kernel_size,
                                 strides=self._strides,
                                 padding="same",
-                                name="{}_seperableconv2d".format(self._name),
+                                name=f"{self._name}_seperableconv2d",
                                 **self._kwargs)(inputs)
-        var_x = Activation("relu", name="{}_relu".format(self._name))(var_x)
+        var_x = Activation("relu", name=f"{self._name}_relu")(var_x)
         return var_x
 
 
@@ -420,7 +430,7 @@ class UpscaleBlock():  # pylint:disable=too-few-public-methods
                  normalization=None,
                  activation="leakyrelu",
                  **kwargs):
-        self._name = _get_name("upscale_{}".format(filters))
+        self._name = _get_name(f"upscale_{filters}")
         logger.debug("name: %s. filters: %s, kernel_size: %s, padding: %s, scale_factor: %s, "
                      "normalization: %s, activation: %s, kwargs: %s)",
                      self._name, filters, kernel_size, padding, scale_factor, normalization,
@@ -453,10 +463,10 @@ class UpscaleBlock():  # pylint:disable=too-few-public-methods
                             padding=self._padding,
                             normalization=self._normalization,
                             activation=self._activation,
-                            name="{}_conv2d".format(self._name),
+                            name=f"{self._name}_conv2d",
                             check_icnr_init=_CONFIG["icnr_init"],
                             **self._kwargs)(inputs)
-        var_x = PixelShuffler(name="{}_pixelshuffler".format(self._name),
+        var_x = PixelShuffler(name=f"{self._name}_pixelshuffler",
                               size=self._scale_factor)(var_x)
         return var_x
 
@@ -501,7 +511,7 @@ class Upscale2xBlock():  # pylint:disable=too-few-public-methods
     """
     def __init__(self, filters, kernel_size=3, padding="same", activation="leakyrelu",
                  interpolation="bilinear", sr_ratio=0.5, scale_factor=2, fast=False, **kwargs):
-        self._name = _get_name("upscale2x_{}_{}".format(filters, "fast" if fast else "hyb"))
+        self._name = _get_name(f"upscale2x_{filters}_{'fast' if fast else 'hyb'}")
 
         self._fast = fast
         self._filters = filters if self._fast else filters - int(filters * sr_ratio)
@@ -536,11 +546,11 @@ class Upscale2xBlock():  # pylint:disable=too-few-public-methods
         if self._fast or (not self._fast and self._filters > 0):
             var_x2 = Conv2D(self._filters, 3,
                             padding=self._padding,
-                            name="{}_conv2d".format(self._name),
+                            name=f"{self._name}_conv2d",
                             **self._kwargs)(var_x)
             var_x2 = UpSampling2D(size=(self._scale_factor, self._scale_factor),
                                   interpolation=self._interpolation,
-                                  name="{}_upsampling2D".format(self._name))(var_x2)
+                                  name=f"{self._name}_upsampling2D")(var_x2)
             if self._fast:
                 var_x1 = UpscaleBlock(self._filters,
                                       kernel_size=self._kernel_size,
@@ -550,7 +560,7 @@ class Upscale2xBlock():  # pylint:disable=too-few-public-methods
                                       **self._kwargs)(var_x)
                 var_x = Add()([var_x2, var_x1])
             else:
-                var_x = Concatenate(name="{}_concatenate".format(self._name))([var_x_sr, var_x2])
+                var_x = Concatenate(name=f"{self._name}_concatenate")([var_x_sr, var_x2])
         else:
             var_x = var_x_sr
         return var_x
@@ -587,7 +597,7 @@ class UpscaleResizeImagesBlock():  # pylint:disable=too-few-public-methods
     """
     def __init__(self, filters, kernel_size=3, padding="same", activation="leakyrelu",
                  scale_factor=2, interpolation="bilinear"):
-        self._name = _get_name("upscale_ri_{}".format(filters))
+        self._name = _get_name(f"upscale_ri_{filters}")
         self._interpolation = interpolation
         self._size = scale_factor
         self._filters = filters
@@ -612,23 +622,23 @@ class UpscaleResizeImagesBlock():  # pylint:disable=too-few-public-methods
 
         var_x_sr = KResizeImages(size=self._size,
                                  interpolation=self._interpolation,
-                                 name="{}_resize".format(self._name))(var_x)
+                                 name=f"{self._name}_resize")(var_x)
         var_x_sr = Conv2D(self._filters, self._kernel_size,
                           strides=1,
                           padding=self._padding,
-                          name="{}_conv".format(self._name))(var_x_sr)
+                          name=f"{self._name}_conv")(var_x_sr)
         var_x_us = Conv2DTranspose(self._filters, 3,
                                    strides=2,
                                    padding=self._padding,
-                                   name="{}_convtrans".format(self._name))(var_x)
+                                   name=f"{self._name}_convtrans")(var_x)
         var_x = Add()([var_x_sr, var_x_us])
 
         if self._activation == "leakyrelu":
-            var_x = LeakyReLU(0.2, name="{}_leakyrelu".format(self._name))(var_x)
+            var_x = LeakyReLU(0.2, name=f"{self._name}_leakyrelu")(var_x)
         if self._activation == "swish":
-            var_x = Swish(name="{}_swish".format(self._name))(var_x)
+            var_x = Swish(name=f"{self._name}_swish")(var_x)
         if self._activation == "prelu":
-            var_x = PReLU(name="{}_prelu".format(self._name))(var_x)
+            var_x = PReLU(name=f"{self._name}_prelu")(var_x)
         return var_x
 
 
@@ -656,7 +666,7 @@ class ResidualBlock():  # pylint:disable=too-few-public-methods
         The output tensor from the Upscale layer
     """
     def __init__(self, filters, kernel_size=3, padding="same", **kwargs):
-        self._name = _get_name("residual_{}".format(filters))
+        self._name = _get_name(f"residual_{filters}")
         logger.debug("name: %s, filters: %s, kernel_size: %s, padding: %s, kwargs: %s)",
                      self._name, filters, kernel_size, padding, kwargs)
         self._use_reflect_padding = _CONFIG["reflect_padding"]
@@ -683,17 +693,17 @@ class ResidualBlock():  # pylint:disable=too-few-public-methods
         if self._use_reflect_padding:
             var_x = ReflectionPadding2D(stride=1,
                                         kernel_size=self._kernel_size,
-                                        name="{}_reflectionpadding2d_0".format(self._name))(var_x)
+                                        name=f"{self._name}_reflectionpadding2d_0")(var_x)
         var_x = Conv2D(self._filters,
                        kernel_size=self._kernel_size,
                        padding=self._padding,
-                       name="{}_conv2d_0".format(self._name),
+                       name=f"{self._name}_conv2d_0",
                        **self._kwargs)(var_x)
-        var_x = LeakyReLU(alpha=0.2, name="{}_leakyrelu_1".format(self._name))(var_x)
+        var_x = LeakyReLU(alpha=0.2, name=f"{self._name}_leakyrelu_1")(var_x)
         if self._use_reflect_padding:
             var_x = ReflectionPadding2D(stride=1,
                                         kernel_size=self._kernel_size,
-                                        name="{}_reflectionpadding2d_1".format(self._name))(var_x)
+                                        name=f"{self._name}_reflectionpadding2d_1")(var_x)
 
         kwargs = {key: val for key, val in self._kwargs.items() if key != "kernel_initializer"}
         if not _CONFIG["conv_aware_init"]:
@@ -703,9 +713,9 @@ class ResidualBlock():  # pylint:disable=too-few-public-methods
         var_x = Conv2D(self._filters,
                        kernel_size=self._kernel_size,
                        padding=self._padding,
-                       name="{}_conv2d_1".format(self._name),
+                       name=f"{self._name}_conv2d_1",
                        **kwargs)(var_x)
 
         var_x = Add()([var_x, inputs])
-        var_x = LeakyReLU(alpha=0.2, name="{}_leakyrelu_3".format(self._name))(var_x)
+        var_x = LeakyReLU(alpha=0.2, name=f"{self._name}_leakyrelu_3")(var_x)
         return var_x

lib/model/normalization/normalization_common.py

@@ -5,22 +5,19 @@ from ast import Import
 import sys
 import inspect
 
-from keras.layers import Layer, InputSpec
-from keras import initializers, regularizers, constraints
-from keras import backend as K
-
-try:
-    from keras.utils import get_custom_objects
-except ImportError:
-    from tensorflow.keras.utils import get_custom_objects
-
 from lib.utils import get_backend
 
-
 if get_backend() == "amd":
-    from keras.backend import normalize_data_format  # pylint:disable=ungrouped-imports
+    from keras.utils import get_custom_objects  # pylint:disable=no-name-in-module
+    from keras.layers import Layer, InputSpec
+    from keras import initializers, regularizers, constraints, backend as K
+    from keras.backend import normalize_data_format  # pylint:disable=no-name-in-module
 else:
-    from tensorflow.python.keras.utils.conv_utils import normalize_data_format
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.utils import get_custom_objects  # noqa pylint:disable=no-name-in-module,import-error
+    from tensorflow.keras.layers import Layer, InputSpec  # noqa pylint:disable=no-name-in-module,import-error
+    from tensorflow.keras import initializers, regularizers, constraints, backend as K  # noqa pylint:disable=no-name-in-module,import-error
+    from tensorflow.python.keras.utils.conv_utils import normalize_data_format  # noqa pylint:disable=no-name-in-module
 
 
 class InstanceNormalization(Layer):
@@ -66,6 +63,7 @@ class InstanceNormalization(Layer):
         - Instance Normalization: The Missing Ingredient for Fast Stylization - \
         https://arxiv.org/abs/1607.08022
     """
+    # pylint:disable=too-many-instance-attributes,too-many-arguments
     def __init__(self,
                  axis=None,
                  epsilon=1e-3,
@@ -353,6 +351,7 @@ class GroupNormalization(Layer):
     ----------
     Shaoanlu GAN: https://github.com/shaoanlu/faceswap-GAN
     """
+    # pylint:disable=too-many-instance-attributes
     def __init__(self, axis=-1, gamma_init='one', beta_init='zero', gamma_regularizer=None,
                  beta_regularizer=None, epsilon=1e-6, group=32, data_format=None, **kwargs):
         self.beta = None

lib/model/normalization/normalization_tf.py

@@ -4,10 +4,10 @@ import inspect
 import sys
 
 import tensorflow as tf
-import tensorflow.keras.backend as K
-# tf.keras has a LayerNormaliztion implementation
-from tensorflow.keras.layers import Layer, LayerNormalization  # noqa pylint:disable=unused-import
-from tensorflow.keras.utils import get_custom_objects
+# Ignore linting errors from Tensorflow's thoroughly broken import system
+from tensorflow.keras import backend as K  # pylint:disable=import-error
+from tensorflow.keras.layers import Layer, LayerNormalization  # noqa pylint:disable=no-name-in-module,unused-import,import-error
+from tensorflow.keras.utils import get_custom_objects  # noqa pylint:disable=no-name-in-module,import-error
 
 
 class RMSNormalization(Layer):
@@ -117,7 +117,8 @@ class RMSNormalization(Layer):
             mean_square = K.mean(K.square(inputs), axis=self.axis, keepdims=True)
         else:
             partial_size = int(layer_size * self.partial)
-            partial_x, _ = tf.split(inputs,
+            partial_x, _ = tf.split(  # pylint:disable=redundant-keyword-arg,no-value-for-parameter
+                inputs,
                 [partial_size, layer_size - partial_size],
                 axis=self.axis)
             mean_square = K.mean(K.square(partial_x), axis=self.axis, keepdims=True)

lib/model/optimizers_plaid.py

@@ -4,7 +4,7 @@ import inspect
 import sys
 
 from keras import backend as K
-from keras.optimizers import Optimizer
+from keras.optimizers import Optimizer, Adam, Nadam, RMSprop  # noqa pylint:disable=unused-import
 from keras.utils import get_custom_objects
 
 

lib/model/optimizers_tf.py

@@ -9,10 +9,9 @@ import sys
 
 import tensorflow as tf
 
-try:
-    from keras.utils import get_custom_objects
-except ImportError:
-    from tensorflow.keras.utils import get_custom_objects
+# Ignore linting errors from Tensorflow's thoroughly broken import system
+from tensorflow.keras.optimizers import (Adam, Nadam, RMSprop)  # noqa pylint:disable=no-name-in-module,unused-import,import-error
+from tensorflow.keras.utils import get_custom_objects  # noqa pylint:disable=no-name-in-module,import-error
 
 
 class AdaBelief(tf.keras.optimizers.Optimizer):
@@ -132,6 +131,7 @@ class AdaBelief(tf.keras.optimizers.Optimizer):
     def __init__(self, learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-14,
                  weight_decay=0.0, rectify=True, amsgrad=False, sma_threshold=5.0, total_steps=0,
                  warmup_proportion=0.1, min_lr=0.0, name="AdaBeliefOptimizer", **kwargs):
+        # pylint:disable=too-many-arguments
         super().__init__(name, **kwargs)
         self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
         self._set_hyper("beta_1", beta_1)
@@ -200,7 +200,7 @@ class AdaBelief(tf.keras.optimizers.Optimizer):
         return wd_t
 
     def _resource_apply_dense(self, grad, handle, apply_state=None):
-        # pylint:disable=too-many-locals
+        # pylint:disable=too-many-locals,unused-argument
         """ Add ops to apply dense gradients to the variable handle.
 
         Parameters
@@ -278,7 +278,7 @@ class AdaBelief(tf.keras.optimizers.Optimizer):
         return tf.group(*updates)
 
     def _resource_apply_sparse(self, grad, handle, indices, apply_state=None):
-        # pylint:disable=too-many-locals
+        # pylint:disable=too-many-locals, unused-argument
         """ Add ops to apply sparse gradients to the variable handle.
 
         Similar to _apply_sparse, the indices argument to this method has been de-duplicated.
@@ -328,7 +328,7 @@ class AdaBelief(tf.keras.optimizers.Optimizer):
         m_corr_t = m_t / (1.0 - beta_1_power)
 
         var_v = self.get_slot(handle, "v")
-        m_t_indices = tf.gather(m_t, indices)
+        m_t_indices = tf.gather(m_t, indices)  # pylint:disable=no-value-for-parameter
         v_scaled_g_values = tf.math.square(grad - m_t_indices) * (1 - beta_2_t)
         v_t = var_v.assign(var_v * beta_2_t + epsilon_t, use_locking=self._use_locking)
         v_t = self._resource_scatter_add(var_v, indices, v_scaled_g_values)
@@ -359,7 +359,9 @@ class AdaBelief(tf.keras.optimizers.Optimizer):
 
         var_update = self._resource_scatter_add(handle,
                                                 indices,
-                                                tf.gather(tf.math.negative(lr_t) * var_t, indices))
+                                                tf.gather(  # pylint:disable=no-value-for-parameter
+                                                    tf.math.negative(lr_t) * var_t,
+                                                    indices))
 
         updates = [var_update, m_t, v_t]
         if self.amsgrad:
@@ -395,6 +397,6 @@ class AdaBelief(tf.keras.optimizers.Optimizer):
 
 
 # Update layers into Keras custom objects
-for name, obj in inspect.getmembers(sys.modules[__name__]):
+for _name, obj in inspect.getmembers(sys.modules[__name__]):
     if inspect.isclass(obj) and obj.__module__ == __name__:
-        get_custom_objects().update({name: obj})
+        get_custom_objects().update({_name: obj})

lib/model/session.py

@@ -5,12 +5,17 @@ import logging
 
 import numpy as np
 import tensorflow as tf
-# pylint:disable=no-name-in-module,import-error
-from keras.layers import Activation
-from keras.models import load_model as k_load_model, Model
 
 from lib.utils import get_backend
 
+if get_backend() == "amd":
+    from keras.layers import Activation
+    from keras.models import load_model as k_load_model, Model
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import Activation  # noqa pylint:disable=no-name-in-module,import-error
+    from tensorflow.keras.models import load_model as k_load_model, Model  # noqa pylint:disable=no-name-in-module,import-error
+
 logger = logging.getLogger(__name__)  # pylint:disable=invalid-name
 
 
@@ -54,7 +59,7 @@ class KSession():
         self._backend = get_backend()
         self._set_session(allow_growth, exclude_gpus)
         self._model_path = model_path
-        self._model_kwargs = dict() if not model_kwargs else model_kwargs
+        self._model_kwargs = {} if not model_kwargs else model_kwargs
         self._model = None
         logger.trace("Initialized: %s", self.__class__.__name__,)
 
@@ -92,7 +97,7 @@ class KSession():
             feed = [feed]
         items = feed[0].shape[0]
         done_items = 0
-        results = list()
+        results = []
         while done_items < items:
             if batch_size < 4:  # Not much difference in BS < 4
                 batch_size = 1

lib/utils.py

@@ -1,7 +1,6 @@
 #!/usr/bin python3
 """ Utilities available across all scripts """
 
-import importlib
 import json
 import logging
 import os
@@ -22,6 +21,7 @@ _image_extensions = [  # pylint:disable=invalid-name
 _video_extensions = [  # pylint:disable=invalid-name
     ".avi", ".flv", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm", ".wmv",
     ".ts", ".vob"]
+_TF_VERS = None
 
 
 class _Backend():  # pylint:disable=too-few-public-methods
@@ -60,8 +60,7 @@ class _Backend():  # pylint:disable=too-few-public-methods
         # Check if environment variable is set, if so use that
         if "FACESWAP_BACKEND" in os.environ:
             fs_backend = os.environ["FACESWAP_BACKEND"].lower()
-            print("Setting Faceswap backend from environment variable to "
-                  "{}".format(fs_backend.upper()))
+            print(f"Setting Faceswap backend from environment variable to {fs_backend.upper()}")
             return fs_backend
         # Intercept for sphinx docs build
         if sys.argv[0].endswith("sphinx-build"):
@@ -70,7 +69,7 @@ class _Backend():  # pylint:disable=too-few-public-methods
             self._configure_backend()
         while True:
             try:
-                with open(self._config_file, "r") as cnf:
+                with open(self._config_file, "r", encoding="utf8") as cnf:
                     config = json.load(cnf)
                 break
             except json.decoder.JSONDecodeError:
@@ -80,7 +79,7 @@ class _Backend():  # pylint:disable=too-few-public-methods
         if fs_backend is None or fs_backend.lower() not in self._backends.values():
             fs_backend = self._configure_backend()
         if current_process().name == "MainProcess":
-            print("Setting Faceswap backend to {}".format(fs_backend.upper()))
+            print(f"Setting Faceswap backend to {fs_backend.upper()}")
         return fs_backend.lower()
 
     def _configure_backend(self):
@@ -95,14 +94,14 @@ class _Backend():  # pylint:disable=too-few-public-methods
         while True:
             selection = input("1: AMD, 2: CPU, 3: NVIDIA, 4: Apple: ")
             if selection not in ("1", "2", "3", "4"):
-                print("'{}' is not a valid selection. Please try again".format(selection))
+                print(f"'{selection}' is not a valid selection. Please try again")
                 continue
             break
         fs_backend = self._backends[selection].lower()
         config = {"backend": fs_backend}
-        with open(self._config_file, "w") as cnf:
+        with open(self._config_file, "w", encoding="utf8") as cnf:
             json.dump(config, cnf)
-        print("Faceswap config written to: {}".format(self._config_file))
+        print(f"Faceswap config written to: {self._config_file}")
         return fs_backend
 
 
@@ -132,6 +131,21 @@ def set_backend(backend):
     _FS_BACKEND = backend.lower()
 
 
+def get_tf_version():
+    """ Obtain the major.minor version of currently installed Tensorflow.
+
+    Returns
+    -------
+    float
+        The currently installed tensorflow version
+    """
+    global _TF_VERS  # pylint:disable=global-statement
+    if _TF_VERS is None:
+        import tensorflow as tf  # pylint:disable=import-outside-toplevel
+        _TF_VERS = float(".".join(tf.__version__.split(".")[:2]))  # pylint:disable=no-member
+    return _TF_VERS
+
+
 def get_folder(path, make_folder=True):
     """ Return a path to a folder, creating it if it doesn't exist
 
@@ -176,7 +190,7 @@ def get_image_paths(directory, extension=None):
     """
     logger = logging.getLogger(__name__)  # pylint:disable=invalid-name
     image_extensions = _image_extensions if extension is None else [extension]
-    dir_contents = list()
+    dir_contents = []
 
     if not os.path.exists(directory):
         logger.debug("Creating folder: '%s'", directory)
@@ -242,7 +256,7 @@ def full_path_split(path):
     >>> ["foo", "baz", "bar"]
     """
     logger = logging.getLogger(__name__)  # pylint:disable=invalid-name
-    allparts = list()
+    allparts = []
     while True:
         parts = os.path.split(path)
         if parts[0] == path:   # sentinel for absolute paths
@@ -297,9 +311,9 @@ def deprecation_warning(function, additional_info=None):
     """
     logger = logging.getLogger(__name__)  # pylint:disable=invalid-name
     logger.debug("func_name: %s, additional_info: %s", function, additional_info)
-    msg = "{}  has been deprecated and will be removed from a future update.".format(function)
+    msg = f"{function}  has been deprecated and will be removed from a future update."
     if additional_info is not None:
-        msg += " {}".format(additional_info)
+        msg += f" {additional_info}"
     logger.warning(msg)
 
 
@@ -355,7 +369,7 @@ class FaceswapError(Exception):
     pass  # pylint:disable=unnecessary-pass
 
 
-class GetModel():  # Pylint:disable=too-few-public-methods
+class GetModel():  # pylint:disable=too-few-public-methods
     """ Check for models in their cache path.
 
     If available, return the path, if not available, get, unzip and install model
@@ -428,7 +442,7 @@ class GetModel():  # Pylint:disable=too-few-public-methods
     @property
     def _model_zip_path(self):
         """ str: The full path to downloaded zip file. """
-        retval = os.path.join(self._cache_dir, "{}.zip".format(self._model_full_name))
+        retval = os.path.join(self._cache_dir, f"{self._model_full_name}.zip")
         self.logger.trace(retval)
         return retval
 
@@ -462,8 +476,8 @@ class GetModel():  # Pylint:disable=too-few-public-methods
     @property
     def _url_download(self):
         """ strL Base download URL for models. """
-        tag = "v{}.{}.{}".format(self._url_section, self._git_model_id, self._model_version)
-        retval = "{}/{}/{}.zip".format(self._url_base, tag, self._model_full_name)
+        tag = f"v{self._url_section}.{self._git_model_id}.{self._model_version}"
+        retval = f"{self._url_base}/{tag}/{self._model_full_name}.zip"
         self.logger.trace("Download url: %s", retval)
         return retval
 
@@ -493,8 +507,8 @@ class GetModel():  # Pylint:disable=too-few-public-methods
                 downloaded_size = self._url_partial_size
                 req = urllib.request.Request(self._url_download)
                 if downloaded_size != 0:
-                    req.add_header("Range", "bytes={}-".format(downloaded_size))
-                response = urllib.request.urlopen(req, timeout=10)
+                    req.add_header("Range", f"bytes={downloaded_size}-")
+                with urllib.request.urlopen(req, timeout=10) as response:
                     self.logger.debug("header info: {%s}", response.info())
                     self.logger.debug("Return Code: %s", response.getcode())
                     self._write_zipfile(response, downloaded_size)
@@ -548,7 +562,7 @@ class GetModel():  # Pylint:disable=too-few-public-methods
         """ Unzip the model file to the cache folder """
         self.logger.info("Extracting: '%s'", self._model_name)
         try:
-            zip_file = zipfile.ZipFile(self._model_zip_path, "r")
+            with zipfile.ZipFile(self._model_zip_path, "r") as zip_file:
                 self._write_model(zip_file)
         except Exception as err:  # pylint:disable=broad-except
             self.logger.error("Unable to extract model file: %s", str(err))
@@ -583,73 +597,3 @@ class GetModel():  # Pylint:disable=too-few-public-methods
                     out_file.write(buffer)
         zip_file.close()
         pbar.close()
-
-
-class KerasFinder(importlib.abc.MetaPathFinder):
-    """ Importlib Abstract Base Class for intercepting the import of Keras and returning either
-    Keras (AMD backend) or tensorflow.keras (any other backend).
-
-    The Importlib documentation is sparse at best, and real world examples are pretty much
-    non-existent. Coupled with this, the import ``tensorflow.keras`` does not resolve so we need
-    to split out to the actual location of Keras within ``tensorflow_core``. This method works, but
-    it relies on hard coded paths, and is likely to not be the most robust.
-
-    A custom loader is not used, as we can use the standard loader once we have returned the
-    correct spec.
-    """
-    def __init__(self):
-        self._logger = logging.getLogger(__name__)
-        self._backend = get_backend()
-        self._tf_keras_locations = [["tensorflow_core", "python", "keras", "api", "_v2"],
-                                    ["tensorflow", "python", "keras", "api", "_v2"]]
-
-    def find_spec(self, fullname, path, target=None):  # pylint:disable=unused-argument
-        """ Obtain the spec for either keras or tensorflow.keras depending on the backend in use.
-
-        If keras is not passed in as part of the :attr:`fullname` or the path is not ``None``
-        (i.e this is a dependency import) then this returns ``None`` to use the standard import
-        library.
-
-        Parameters
-        ----------
-        fullname: str
-            The absolute name of the module to be imported
-        path: str
-            The search path for the module
-        target: module object, optional
-            Inherited from parent but unused
-
-        Returns
-        -------
-        :class:`importlib.ModuleSpec`
-            The spec for the Keras module to be imported
-        """
-        prefix = fullname.split(".")[0]
-        suffix = fullname.split(".")[-1]
-        if prefix != "keras" or path is not None:
-            return None
-        self._logger.debug("Importing '%s' as keras for backend: '%s'",
-                           "keras" if self._backend == "amd" else "tf.keras", self._backend)
-        path = sys.path if path is None else path
-        for entry in path:
-            locations = ([os.path.join(entry, *location)
-                          for location in self._tf_keras_locations]
-                         if self._backend != "amd" else [entry])
-            for location in locations:
-                self._logger.debug("Scanning: '%s' for '%s'", location, suffix)
-                if os.path.isdir(os.path.join(location, suffix)):
-                    filename = os.path.join(location, suffix, "__init__.py")
-                    submodule_locations = [os.path.join(location, suffix)]
-                else:
-                    filename = os.path.join(location, suffix + ".py")
-                    submodule_locations = None
-                if not os.path.exists(filename):
-                    continue
-                retval = importlib.util.spec_from_file_location(
-                    fullname,
-                    filename,
-                    submodule_search_locations=submodule_locations)
-                self._logger.debug("Found spec: %s", retval)
-                return retval
-        self._logger.debug("Spec not found for '%s'. Falling back to default import", fullname)
-        return None

locales/es/LC_MESSAGES/tools.sort.cli.po

@@ -184,7 +184,7 @@ msgid ""
 "right. If the number of images doesn't divide evenly into the number of "
 "bins, the remaining images get put in the last bin. For black-pixels it "
 "represents the divider of the percentage of black pixels. For 10, first "
-"folder will have the faces with 0 to 10% black pixels, second 11 to 20%, "
+"folder will have the faces with 0 to 10%% black pixels, second 11 to 20%%, "
 "etc. Default value: 5"
 msgstr ""
 "Valor entero. Número de carpetas que se utilizarán al agrupar por 'blur' y "
@@ -196,8 +196,8 @@ msgstr ""
 "caras que miren más a la derecha. Si el número de imágenes no se divide "
 "uniformemente en el número de carpetas, las imágenes restantes se colocan en "
 "la última carpeta. Para píxeles negros, representa el divisor del porcentaje "
-"de píxeles negros. Para 10, la primera carpeta tendrá las caras con 0 a 10% "
-"de píxeles negros, la segunda de 11 a 20%, etc. Valor por defecto: 5"
+"de píxeles negros. Para 10, la primera carpeta tendrá las caras con 0 a 10%% "
+"de píxeles negros, la segunda de 11 a 20%%, etc. Valor por defecto: 5"
 
 #: tools/sort/cli.py:154 tools/sort/cli.py:164
 msgid "settings"

locales/tools.sort.cli.pot

@@ -85,7 +85,7 @@ msgid "Group by method. When -fp/--final-processing by folders choose the how th
 msgstr ""
 
 #: tools/sort/cli.py:140
-msgid "Integer value. Number of folders that will be used to group by blur, face-yaw and black-pixels. For blur folder 0 will be the least blurry, while the last folder will be the blurriest. For face-yaw the number of bins is by how much 180 degrees is divided. So if you use 18, then each folder will be a 10 degree increment. Folder 0 will contain faces looking the most to the left whereas the last folder will contain the faces looking the most to the right. If the number of images doesn't divide evenly into the number of bins, the remaining images get put in the last bin. For black-pixels it represents the divider of the percentage of black pixels. For 10, first folder will have the faces with 0 to 10% black pixels, second 11 to 20%, etc. Default value: 5"
+msgid "Integer value. Number of folders that will be used to group by blur, face-yaw and black-pixels. For blur folder 0 will be the least blurry, while the last folder will be the blurriest. For face-yaw the number of bins is by how much 180 degrees is divided. So if you use 18, then each folder will be a 10 degree increment. Folder 0 will contain faces looking the most to the left whereas the last folder will contain the faces looking the most to the right. If the number of images doesn't divide evenly into the number of bins, the remaining images get put in the last bin. For black-pixels it represents the divider of the percentage of black pixels. For 10, first folder will have the faces with 0 to 10%% black pixels, second 11 to 20%%, etc. Default value: 5"
 msgstr ""
 
 #: tools/sort/cli.py:154 tools/sort/cli.py:164

plugins/extract/detect/mtcnn.py

@@ -5,11 +5,17 @@ from __future__ import absolute_import, division, print_function
 
 import cv2
 import numpy as np
-# pylint:disable=import-error
-from keras.layers import Conv2D, Dense, Flatten, Input, MaxPool2D, Permute, PReLU
+
 from lib.model.session import KSession
+from lib.utils import get_backend
 from ._base import Detector, logger
 
+if get_backend() == "amd":
+    from keras.layers import Conv2D, Dense, Flatten, Input, MaxPool2D, Permute, PReLU
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import Conv2D, Dense, Flatten, Input, MaxPool2D, Permute, PReLU  # noqa pylint:disable=no-name-in-module,import-error
+
 
 class Detect(Detector):
     """ MTCNN detector for face recognition """
@@ -234,8 +240,8 @@ class MTCNN():
         rectangles = self.detect_pnet(batch, origin_h, origin_w)
         rectangles = self.detect_rnet(batch, rectangles, origin_h, origin_w)
         rectangles = self.detect_onet(batch, rectangles, origin_h, origin_w)
-        ret_boxes = list()
-        ret_points = list()
+        ret_boxes = []
+        ret_points = []
         for rects in rectangles:
             if rects:
                 total_boxes = np.array([result[:5] for result in rects])
@@ -284,7 +290,7 @@ class MTCNN():
         # TODO: batching
         for idx, rectangles in enumerate(rectangle_batch):
             if not rectangles:
-                ret.append(list())
+                ret.append([])
                 continue
             image = images[idx]
             crop_number = 0
@@ -307,11 +313,11 @@ class MTCNN():
 
     def detect_onet(self, images, rectangle_batch, height, width):
         """ third stage - further refinement and facial landmarks positions with o-net """
-        ret = list()
+        ret = []
         # TODO: batching
         for idx, rectangles in enumerate(rectangle_batch):
             if not rectangles:
-                ret.append(list())
+                ret.append([])
                 continue
             image = images[idx]
             crop_number = 0

plugins/extract/detect/s3fd.py

@@ -8,13 +8,22 @@ https://github.com/1adrianb/face-alignment
 
 from scipy.special import logsumexp
 import numpy as np
-import keras  # pylint:disable=import-error
-import keras.backend as K  # pylint:disable=import-error
-from keras.layers import Concatenate, Conv2D, Input, Maximum, MaxPooling2D, ZeroPadding2D
 
 from lib.model.session import KSession
+from lib.utils import get_backend
 from ._base import Detector, logger
 
+if get_backend() == "amd":
+    import keras
+    from keras import backend as K
+    from keras.layers import Concatenate, Conv2D, Input, Maximum, MaxPooling2D, ZeroPadding2D
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow import keras
+    from tensorflow.keras import backend as K  # pylint:disable=import-error
+    from tensorflow.keras.layers import (  # pylint:disable=no-name-in-module,import-error
+        Concatenate, Conv2D, Input, Maximum, MaxPooling2D, ZeroPadding2D)
+
 
 class Detect(Detector):
     """ S3FD detector for face recognition """
@@ -316,11 +325,11 @@ class S3fd(KSession):
         tensor
             The output tensor from the convolution block
         """
-        name = "conv{}".format(idx)
+        name = f"conv{idx}"
         var_x = inputs
         for i in range(1, recursions + 1):
-            rec_name = "{}_{}".format(name, i)
-            var_x = ZeroPadding2D(1, name="{}.zeropad".format(rec_name))(var_x)
+            rec_name = f"{name}_{i}"
+            var_x = ZeroPadding2D(1, name=f"{rec_name}.zeropad")(var_x)
             var_x = Conv2D(filters,
                            kernel_size=3,
                            strides=1,
@@ -346,13 +355,13 @@ class S3fd(KSession):
         tensor
             The output tensor from the convolution block
         """
-        name = "conv{}".format(idx)
+        name = f"conv{idx}"
         var_x = inputs
         for i in range(1, 3):
-            rec_name = "{}_{}".format(name, i)
+            rec_name = f"{name}_{i}"
             size = 1 if i == 1 else 3
             if i == 2:
-                var_x = ZeroPadding2D(1, name="{}.zeropad".format(rec_name))(var_x)
+                var_x = ZeroPadding2D(1, name=f"{rec_name}.zeropad")(var_x)
             var_x = Conv2D(filters * i,
                            kernel_size=size,
                            strides=i,
@@ -386,7 +395,7 @@ class S3fd(KSession):
         bounding_boxes_scales: list
             The output predictions from the S3FD model
         """
-        ret = list()
+        ret = []
         batch_size = range(bounding_boxes_scales[0].shape[0])
         for img in batch_size:
             bboxlist = [scale[img:img+1] for scale in bounding_boxes_scales]
@@ -399,7 +408,7 @@ class S3fd(KSession):
         """ Perform post processing on output
             TODO: do this on the batch.
         """
-        retval = list()
+        retval = []
         for i in range(len(bboxlist) // 2):
             bboxlist[i * 2] = self.softmax(bboxlist[i * 2], axis=3)
         for i in range(len(bboxlist) // 2):
@@ -450,7 +459,7 @@ class S3fd(KSession):
     @staticmethod
     def _nms(boxes, threshold):
         """ Perform Non-Maximum Suppression """
-        retained_box_indices = list()
+        retained_box_indices = []
 
         areas = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
         ranked_indices = boxes[:, 4].argsort()[::-1]

plugins/extract/mask/bisenet_fp.py

@@ -4,24 +4,35 @@
 Architecture and Pre-Trained Model ported from PyTorch to Keras by TorzDF from
 https://github.com/zllrunning/face-parsing.PyTorch
 """
-
 import numpy as np
 
-from keras import backend as K
-from keras.layers import (Activation, Add, BatchNormalization, Concatenate, Conv2D,
-                          GlobalAveragePooling2D, Input, MaxPooling2D, Multiply, Reshape,
-                          UpSampling2D, ZeroPadding2D)
-
 from lib.model.session import KSession
+from lib.utils import get_backend
+from plugins.extract._base import _get_config
 from ._base import Masker, logger
 
+if get_backend() == "amd":
+    from keras import backend as K
+    from keras.layers import (
+        Activation, Add, BatchNormalization, Concatenate, Conv2D, GlobalAveragePooling2D, Input,
+        MaxPooling2D, Multiply, Reshape, UpSampling2D, ZeroPadding2D)
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import backend as K  # pylint:disable=import-error
+    from tensorflow.keras.layers import (  # pylint:disable=no-name-in-module,import-error
+        Activation, Add, BatchNormalization, Concatenate, Conv2D, GlobalAveragePooling2D, Input,
+        MaxPooling2D, Multiply, Reshape, UpSampling2D, ZeroPadding2D)
+
 
 class Mask(Masker):
     """ Neural network to process face image into a segmentation mask of the face """
     def __init__(self, **kwargs):
+        self._is_faceswap = self._check_weights_selection(kwargs.get("configfile"))
+
         git_model_id = 14
-        model_filename = "bisnet_face_parsing_v1.h5"
+        model_filename = f"bisnet_face_parsing_v{'2' if self._is_faceswap else '1'}.h5"
         super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs)
+
         self.name = "BiSeNet - Face Parsing"
         self.input_size = 512
         self.color_format = "RGB"
@@ -29,12 +40,33 @@ class Mask(Masker):
         self.vram_warnings = 256
         self.vram_per_batch = 64
         self.batchsize = self.config["batch-size"]
-        self._segment_indices = self._get_segment_indices()
 
+        self._segment_indices = self._get_segment_indices()
         self._storage_centering = "head" if self.config["include_hair"] else "face"
         # Separate storage for face and head masks
         self._storage_name = f"{self._storage_name}_{self._storage_centering}"
 
+    def _check_weights_selection(self, configfile):
+        """ Check which weights have been selected.
+
+        This is required for passing along the correct file name for the corresponding weights
+        selection, so config needs to be loaded and scanned prior to parent loading it.
+
+        Parameters
+        ----------
+        configfile: str
+            Path to a custom configuration ``ini`` file. ``None`` to use system configfile
+
+        Returns
+        -------
+        bool
+            ``True`` if `faceswap` trained weights have been selected. ``False`` if `original`
+            weights have been selected
+        """
+        config = _get_config(".".join(self.__module__.split(".")[-2:]), configfile=configfile)
+        retval = config.get("weights", "faceswap").lower() == "faceswap"
+        return retval
+
     def _get_segment_indices(self):
         """ Obtain the segment indices to include within the face mask area based on user
         configuration settings.
@@ -46,28 +78,33 @@ class Mask(Masker):
 
         Notes
         -----
-        Model segment indices:
+        'original' Model segment indices:
         0: background, 1: skin, 2: left brow, 3: right brow, 4: left eye, 5: right eye, 6: glasses
         7: left ear, 8: right ear, 9: earing, 10: nose, 11: mouth, 12: upper lip, 13: lower_lip,
         14: neck, 15: neck ?, 16: cloth, 17: hair, 18: hat
+
+        'faceswap' Model segment indices:
+        0: background, 1: skin, 2: ears, 3: hair, 4: glasses
         """
-        retval = [1, 2, 3, 4, 5, 10, 11, 12, 13]
+        retval = [1] if self._is_faceswap else [1, 2, 3, 4, 5, 10, 11, 12, 13]
+
         if self.config["include_glasses"]:
-            retval.append(6)
+            retval.append(4 if self._is_faceswap else 6)
         if self.config["include_ears"]:
-            retval.extend([7, 8, 9])
+            retval.extend([2] if self._is_faceswap else [7, 8, 9])
         if self.config["include_hair"]:
-            retval.append(17)
+            retval.append(3 if self._is_faceswap else 17)
         logger.debug("Selected segment indices: %s", retval)
         return retval
 
     def init_model(self):
         """ Initialize the BiSeNet Face Parsing model. """
+        lbls = 5 if self._is_faceswap else 19
         self.model = BiSeNet(self.model_path,
                              self.config["allow_growth"],
                              self._exclude_gpus,
                              self.input_size,
-                             19)
+                             lbls)
 
         placeholder = np.zeros((self.batchsize, self.input_size, self.input_size, 3),
                                dtype="float32")
@@ -75,8 +112,8 @@ class Mask(Masker):
 
     def process_input(self, batch):
         """ Compile the detected faces for prediction """
-        mean = (0.485, 0.456, 0.406)
-        std = (0.229, 0.224, 0.225)
+        mean = (0.384, 0.314, 0.279) if self._is_faceswap else (0.485, 0.456, 0.406)
+        std = (0.324, 0.286, 0.275) if self._is_faceswap else (0.229, 0.224, 0.225)
 
         batch["feed"] = ((np.array([feed.face[..., :3]
                                     for feed in batch["feed_faces"]],

plugins/extract/mask/bisenet_fp_defaults.py

@@ -63,6 +63,18 @@ _DEFAULTS = {
         group="settings",
         gui_radio=False,
         fixed=True),
+    "weights": dict(
+        default="faceswap",
+        info="The trained weights to use.\n"
+             "\n\tfaceswap - Weights trained on wildly varied Faceswap extracted data to better "
+             "handle varying conditions, obstructions, glasses and multiple targets within a "
+             "single extracted image."
+             "\n\toriginal - The original weights trained on the CelebAMask-HQ dataset.",
+        choices=["faceswap", "original"],
+        datatype=str,
+        group="settings",
+        gui_radio=True,
+    ),
     "include_ears": dict(
         default=False,
         info="Whether to include ears within the face mask.",
@@ -77,8 +89,10 @@ _DEFAULTS = {
     ),
     "include_glasses": dict(
         default=True,
-        info="Whether to include glasses within the face mask. NB: excluding glasses will mask "
-             "out the lenses as well as the frames.",
+        info="Whether to include glasses within the face mask.\n\tFor 'original' weights "
+             "excluding glasses will mask out the lenses as well as the frames.\n\tFor 'faceswap' "
+             "weights, the model has been trained to mask out lenses if eyes cannot be seen (i.e. "
+             "dark sunglasses) or just the frames if the eyes can be seen. ",
         datatype=bool,
         group="settings"
     ),

plugins/extract/mask/vgg_clear.py

@@ -2,13 +2,21 @@
 """ VGG Clear face mask plugin. """
 
 import numpy as np
-from keras.layers import (Add, Conv2D,    # pylint:disable=no-name-in-module,import-error
-                          Conv2DTranspose, Cropping2D, Dropout, Input, Lambda,
-                          MaxPooling2D, ZeroPadding2D)
 
 from lib.model.session import KSession
+from lib.utils import get_backend
 from ._base import Masker, logger
 
+if get_backend() == "amd":
+    from keras.layers import (
+        Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D,
+        ZeroPadding2D)
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import (  # pylint:disable=no-name-in-module,import-error
+        Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D,
+        ZeroPadding2D)
+
 
 class Mask(Masker):
     """ Neural network to process face image into a segmentation mask of the face """
@@ -151,7 +159,7 @@ class _ConvBlock():  # pylint:disable=too-few-public-methods
         The number of consecutive Conv2D layers to create
     """
     def __init__(self, level, filters, iterations):
-        self._name = "conv{}_".format(level)
+        self._name = f"conv{level}_"
         self._level = level
         self._filters = filters
         self._iterator = range(1, iterations + 1)
@@ -176,10 +184,10 @@ class _ConvBlock():  # pylint:disable=too-few-public-methods
                            3,
                            padding=padding,
                            activation="relu",
-                           name="{}{}".format(self._name, i))(var_x)
+                           name=f"{self._name}{i}")(var_x)
         var_x = MaxPooling2D(padding="same",
                              strides=(2, 2),
-                             name="pool{}".format(self._level))(var_x)
+                             name=f"pool{self._level}")(var_x)
         return var_x
 
 
@@ -196,7 +204,7 @@ class _ScorePool():  # pylint:disable=too-few-public-methods
         The amount of 2D cropping to apply. Tuple of `ints`
     """
     def __init__(self, level, scale, crop):
-        self._name = "_pool{}".format(level)
+        self._name = f"_pool{level}"
         self._cropping = (crop, crop)
         self._scale = scale
 

plugins/extract/mask/vgg_obstructed.py

@@ -2,13 +2,22 @@
 """ VGG Obstructed face mask plugin """
 
 import numpy as np
-from keras.layers import (Add, Conv2D,  # pylint:disable=no-name-in-module,import-error
-                          Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D,
-                          ZeroPadding2D)
+
 
 from lib.model.session import KSession
+from lib.utils import get_backend
 from ._base import Masker, logger
 
+if get_backend() == "amd":
+    from keras.layers import (
+        Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D,
+        ZeroPadding2D)
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import (  # pylint:disable=no-name-in-module,import-error
+        Add, Conv2D, Conv2DTranspose, Cropping2D, Dropout, Input, Lambda, MaxPooling2D,
+        ZeroPadding2D)
+
 
 class Mask(Masker):
     """ Neural network to process face image into a segmentation mask of the face """
@@ -150,7 +159,7 @@ class _ConvBlock():  # pylint:disable=too-few-public-methods
         The number of consecutive Conv2D layers to create
     """
     def __init__(self, level, filters, iterations):
-        self._name = "conv{}_".format(level)
+        self._name = f"conv{level}_"
         self._level = level
         self._filters = filters
         self._iterator = range(1, iterations + 1)
@@ -175,10 +184,10 @@ class _ConvBlock():  # pylint:disable=too-few-public-methods
                            3,
                            padding=padding,
                            activation="relu",
-                           name="{}{}".format(self._name, i))(var_x)
+                           name=f"{self._name}{i}")(var_x)
         var_x = MaxPooling2D(padding="same",
                              strides=(2, 2),
-                             name="pool{}".format(self._level))(var_x)
+                             name=f"pool{self._level}")(var_x)
         return var_x
 
 
@@ -195,7 +204,7 @@ class _ScorePool():  # pylint:disable=too-few-public-methods
         The amount of 2D cropping to apply
     """
     def __init__(self, level, scale, crop):
-        self._name = "_pool{}".format(level)
+        self._name = f"_pool{level}"
         self._cropping = ((crop, crop), (crop, crop))
         self._scale = scale
 

plugins/train/_config.py

@@ -234,6 +234,7 @@ class Config(FaceswapConfig):
         L_inf_norm https://medium.com/@montjoile/l0-norm-l1-norm-l2-norm-l-infinity
                    -norm-7a7d18a4f40c
         SSIM http://www.cns.nyu.edu/pub/eero/wang03-reprint.pdf
+        MSSIM https://www.cns.nyu.edu/pub/eero/wang03b.pdf
         GMSD https://arxiv.org/ftp/arxiv/papers/1308/1308.3052.pdf
         """
         logger.debug("Setting Loss config")
@@ -248,8 +249,8 @@ class Config(FaceswapConfig):
             datatype=str,
             group="loss",
             default="ssim",
-            choices=["mae", "mse", "logcosh", "smooth_loss", "l_inf_norm", "ssim", "gmsd",
-                     "pixel_gradient_diff"],
+            choices=["mae", "mse", "logcosh", "smooth_loss", "l_inf_norm", "ssim", "ms_ssim",
+                     "gmsd", "pixel_gradient_diff"],
             info="The loss function to use."
                  "\n\t MAE - Mean absolute error will guide reconstructions of each pixel "
                  "towards its median value in the training dataset. Robust to outliers but as "
@@ -269,6 +270,9 @@ class Config(FaceswapConfig):
                  "\n\t SSIM - Structural Similarity Index Metric is a perception-based "
                  "loss that considers changes in texture, luminance, contrast, and local spatial "
                  "statistics of an image. Potentially delivers more realistic looking images."
+                 "\n\t MS_SSIM - Multiscale Structural Similarity Index Metric is similar to SSIM "
+                 "except that it performs the calculations along multiple scales of the input "
+                 "image. NB: This loss currently does not work on AMD Cards."
                  "\n\t GMSD - Gradient Magnitude Similarity Deviation seeks to match "
                  "the global standard deviation of the pixel to pixel differences between two "
                  "images. Similar in approach to SSIM. NB: This loss does not currently work on "
@@ -304,9 +308,10 @@ class Config(FaceswapConfig):
                  "loss functions.\n\nNB: You should only adjust this if you know what you are "
                  "doing!\n\n"
                  "L2 regularization applies a penalty term to the given Loss function. This "
-                 "penalty will only be applied if SSIM or GMSD is selected for the main loss "
-                 "function, otherwise it is ignored.\n\nThe value given here is as a percentage "
-                 "weight of the main loss function. For example:"
+                 "penalty will only be applied if SSIM, MS-SSIM or GMSD is selected for the main "
+                 "loss function, otherwise it is ignored."
+                 "\n\nThe value given here is as a percentage weight of the main loss function. "
+                 "For example:"
                  "\n\t 100 - Will give equal weighting to the main loss and the penalty function. "
                  "\n\t 25 - Will give the penalty function 1/4 of the weight of the main loss "
                  "function. "

plugins/train/model/_base.py

@@ -16,23 +16,26 @@ from contextlib import nullcontext
 import numpy as np
 import tensorflow as tf
 
-from keras import losses as k_losses
-from keras import backend as K
-from keras.layers import Input
-from keras.models import load_model, Model as KModel
-
-try:
-    from keras.optimizers import Adam, Nadam, RMSprop
-except ImportError:
-    from tensorflow.keras.optimizers import Adam, Nadam, RMSprop
-
 from lib.serializer import get_serializer
 from lib.model.backup_restore import Backup
 from lib.model import losses, optimizers
 from lib.model.nn_blocks import set_config as set_nnblock_config
-from lib.utils import get_backend, FaceswapError
+from lib.utils import get_backend, get_tf_version, FaceswapError
 from plugins.train._config import Config
 
+if get_backend() == "amd":
+    from keras import losses as k_losses
+    from keras import backend as K
+    from keras.layers import Input
+    from keras.models import load_model, Model as KModel
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import losses as k_losses  # pylint:disable=import-error
+    from tensorflow.keras import backend as K  # pylint:disable=import-error
+    from tensorflow.keras.layers import Input  # pylint:disable=import-error,no-name-in-module
+    from tensorflow.keras.models import load_model, Model as KModel  # noqa pylint:disable=import-error,no-name-in-module
+
+
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 _CONFIG = None
 
@@ -268,13 +271,14 @@ class ModelBase():
             return
 
         if len(multiple_models) == 1:
-            msg = ("You have requested to train with the '{}' plugin, but a model file for the "
-                   "'{}' plugin already exists in the folder '{}'.\nPlease select a different "
-                   "model folder.".format(self.name, multiple_models[0], self.model_dir))
+            msg = (f"You have requested to train with the '{self.name}' plugin, but a model file "
+                   f"for the '{multiple_models[0]}' plugin already exists in the folder "
+                   f"'{self.model_dir}'.\nPlease select a different model folder.")
         else:
-            msg = ("There are multiple plugin types ('{}') stored in the model folder '{}'. This "
-                   "is not supported.\nPlease split the model files into their own folders before "
-                   "proceeding".format("', '".join(multiple_models), self.model_dir))
+            ptypes = "', '".join(multiple_models)
+            msg = (f"There are multiple plugin types ('{ptypes}') stored in the model folder '"
+                   f"{self.model_dir}'. This is not supported.\nPlease split the model files into "
+                   "their own folders before proceeding")
         raise FaceswapError(msg)
 
     def build(self):
@@ -315,11 +319,11 @@ class ModelBase():
         if not all(os.path.isfile(os.path.join(self.model_dir, fname))
                    for fname in self._legacy_mapping()):
             return
-        archive_dir = "{}_TF1_Archived".format(self.model_dir)
+        archive_dir = f"{self.model_dir}_TF1_Archived"
         if os.path.exists(archive_dir):
             raise FaceswapError("We need to update your model files for use with Tensorflow 2.x, "
                                 "but the archive folder already exists. Please remove the "
-                                "following folder to continue: '{}'".format(archive_dir))
+                                f"following folder to continue: '{archive_dir}'")
 
         logger.info("Updating legacy models for Tensorflow 2.x")
         logger.info("Your Tensorflow 1.x models will be archived in the following location: '%s'",
@@ -375,7 +379,7 @@ class ModelBase():
             input_shapes = [self.input_shape, self.input_shape]
         else:
             input_shapes = self.input_shape
-        inputs = [Input(shape=shape, name="face_in_{}".format(side))
+        inputs = [Input(shape=shape, name=f"face_in_{side}")
                   for side, shape in zip(("a", "b"), input_shapes)]
         logger.debug("inputs: %s", inputs)
         return inputs
@@ -454,7 +458,7 @@ class ModelBase():
         seen = {name: 0 for name in set(self._model.output_names)}
         new_names = []
         for name in self._model.output_names:
-            new_names.append("{}_{}".format(name, seen[name]))
+            new_names.append(f"{name}_{seen[name]}")
             seen[name] += 1
         logger.debug("Output names rewritten: (old: %s, new: %s)",
                      self._model.output_names, new_names)
@@ -514,7 +518,7 @@ class _IO():
     @property
     def _filename(self):
         """str: The filename for this model."""
-        return os.path.join(self._model_dir, "{}.h5".format(self._plugin.name))
+        return os.path.join(self._model_dir, f"{self._plugin.name}.h5")
 
     @property
     def model_exists(self):
@@ -571,7 +575,7 @@ class _IO():
                        "You can try to load the model again but if the problem persists you "
                        "should use the Restore Tool to restore your model from backup.\n"
                        f"Original error: {str(err)}")
-                raise FaceswapError(msg)
+                raise FaceswapError(msg) from err
             raise err
         except KeyError as err:
             if "unable to open object" in str(err).lower():
@@ -580,7 +584,7 @@ class _IO():
                        "You can try to load the model again but if the problem persists you "
                        "should use the Restore Tool to restore your model from backup.\n"
                        f"Original error: {str(err)}")
-                raise FaceswapError(msg)
+                raise FaceswapError(msg) from err
             raise err
 
         logger.info("Loaded model from disk: '%s'", self._filename)
@@ -609,9 +613,9 @@ class _IO():
 
         msg = "[Saved models]"
         if save_averages:
-            lossmsg = ["face_{}: {:.5f}".format(side, avg)
+            lossmsg = [f"face_{side}: {avg:.5f}"
                        for side, avg in zip(("a", "b"), save_averages)]
-            msg += " - Average loss since last save: {}".format(", ".join(lossmsg))
+            msg += f" - Average loss since last save: {', '.join(lossmsg)}"
         logger.info(msg)
 
     def _get_save_averages(self):
@@ -703,12 +707,11 @@ class _Settings():
         logger.debug("Initializing %s: (arguments: %s, mixed_precision: %s, allow_growth: %s, "
                      "is_predict: %s)", self.__class__.__name__, arguments, mixed_precision,
                      allow_growth, is_predict)
-        self._tf_version = [int(i) for i in tf.__version__.split(".")[:2]]
         self._set_tf_settings(allow_growth, arguments.exclude_gpus)
 
         use_mixed_precision = not is_predict and mixed_precision and get_backend() == "nvidia"
         # Mixed precision moved out of experimental in tensorflow 2.4
-        if use_mixed_precision and self._tf_version[0] == 2 and self._tf_version[1] < 4:
+        if use_mixed_precision and get_tf_version() < 2.4:
             self._mixed_precision = tf.keras.mixed_precision.experimental
         elif use_mixed_precision:
             self._mixed_precision = tf.keras.mixed_precision
@@ -747,9 +750,8 @@ class _Settings():
         """
         # tensorflow versions < 2.4 had different kwargs where scaling needs to be explicitly
         # defined
-        vers = self._tf_version
-        kwargs = dict(loss_scale="dynamic") if vers[0] == 2 and vers[1] < 4 else dict()
-        logger.debug("tf version: %s, kwargs: %s", vers, kwargs)
+        kwargs = dict(loss_scale="dynamic") if get_tf_version() < 2.4 else {}
+        logger.debug("tf version: %s, kwargs: %s", get_tf_version(), kwargs)
         return self._mixed_precision.LossScaleOptimizer(optimizer, **kwargs)
 
     @classmethod
@@ -825,10 +827,11 @@ class _Settings():
             return False
         logger.info("Enabling Mixed Precision Training.")
 
-        if exclude_gpus and self._tf_version[0] == 2 and self._tf_version[1] == 2:
+        if exclude_gpus and get_tf_version() == 2.2:
             # TODO remove this hacky fix to disable mixed precision compatibility testing when
             # tensorflow 2.2 support dropped
-            # pylint:disable=import-outside-toplevel,protected-access,import-error
+            # pylint:disable=import-outside-toplevel,protected-access
+            # pylint:disable=import-error,no-name-in-module
             from tensorflow.python.keras.mixed_precision.experimental import \
                 device_compatibility_check
             logger.debug("Overriding tensorflow _logged_compatibility_check parameter. Initial "
@@ -837,7 +840,7 @@ class _Settings():
             logger.debug("New value: %s", device_compatibility_check._logged_compatibility_check)
 
         policy = self._mixed_precision.Policy('mixed_float16')
-        if self._tf_version[0] == 2 and self._tf_version[1] < 4:
+        if get_tf_version() < 2.4:
             self._mixed_precision.set_policy(policy)
         else:
             self._mixed_precision.set_global_policy(policy)
@@ -1106,9 +1109,11 @@ class _Optimizer():  # pylint:disable=too-few-public-methods
                      optimizer, learning_rate, clipnorm, epsilon, arguments)
         valid_optimizers = {"adabelief": (optimizers.AdaBelief,
                                           dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)),
-                            "adam": (Adam, dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)),
-                            "nadam": (Nadam, dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)),
-                            "rms-prop": (RMSprop, dict(epsilon=epsilon))}
+                            "adam": (optimizers.Adam,
+                                     dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)),
+                            "nadam": (optimizers.Nadam,
+                                      dict(beta_1=0.5, beta_2=0.99, epsilon=epsilon)),
+                            "rms-prop": (optimizers.RMSprop, dict(epsilon=epsilon))}
         self._optimizer, self._kwargs = valid_optimizers[optimizer]
 
         self._configure(learning_rate, clipnorm, arguments)
@@ -1172,12 +1177,13 @@ class _Loss():
                                smooth_loss=losses.GeneralizedLoss(),
                                l_inf_norm=losses.LInfNorm(),
                                ssim=losses.DSSIMObjective(),
+                               ms_ssim=losses.MSSSIMLoss(),
                                gmsd=losses.GMSDLoss(),
                                pixel_gradient_diff=losses.GradientLoss())
-        self._uses_l2_reg = ["ssim", "gmsd"]
+        self._uses_l2_reg = ["ssim", "ms_ssim", "gmsd"]
         self._inputs = None
         self._names = []
-        self._funcs = dict()
+        self._funcs = {}
         logger.debug("Initialized: %s", self.__class__.__name__)
 
     @property
@@ -1252,7 +1258,7 @@ class _Loss():
                          side, output_names, output_shapes, output_types)
             self._names.extend(["{}_{}{}".format(name, side,
                                                  "" if output_types.count(name) == 1
-                                                 else "_{}".format(idx))
+                                                 else f"_{idx}")
                                 for idx, name in enumerate(output_types)])
         logger.debug(self._names)
 
@@ -1358,13 +1364,13 @@ class State():
                      "config_changeable_items: '%s', no_logs: %s", self.__class__.__name__,
                      model_dir, model_name, config_changeable_items, no_logs)
         self._serializer = get_serializer("json")
-        filename = "{}_state.{}".format(model_name, self._serializer.file_extension)
+        filename = f"{model_name}_state.{self._serializer.file_extension}"
         self._filename = os.path.join(model_dir, filename)
         self._name = model_name
         self._iterations = 0
-        self._sessions = dict()
-        self._lowest_avg_loss = dict()
-        self._config = dict()
+        self._sessions = {}
+        self._lowest_avg_loss = {}
+        self._config = {}
         self._load(config_changeable_items)
         self._session_id = self._new_session_id()
         self._create_new_session(no_logs, config_changeable_items)
@@ -1477,10 +1483,10 @@ class State():
             return
         state = self._serializer.load(self._filename)
         self._name = state.get("name", self._name)
-        self._sessions = state.get("sessions", dict())
-        self._lowest_avg_loss = state.get("lowest_avg_loss", dict())
+        self._sessions = state.get("sessions", {})
+        self._lowest_avg_loss = state.get("lowest_avg_loss", {})
         self._iterations = state.get("iterations", 0)
-        self._config = state.get("config", dict())
+        self._config = state.get("config", {})
         logger.debug("Loaded state: %s", state)
         self._replace_config(config_changeable_items)
 
@@ -1674,7 +1680,7 @@ class _Inference():  # pylint:disable=too-few-public-methods
         logger.debug("Compiling inference model. saved_model: %s", saved_model)
         struct = self._get_filtered_structure()
         model_inputs = self._get_inputs(saved_model.inputs)
-        compiled_layers = dict()
+        compiled_layers = {}
         for layer in saved_model.layers:
             if layer.name not in struct:
                 logger.debug("Skipping unused layer: '%s'", layer.name)
@@ -1708,7 +1714,7 @@ class _Inference():  # pylint:disable=too-few-public-methods
                 logger.debug("Compiling layer '%s': layer inputs: %s", layer.name, layer_inputs)
                 model = layer(layer_inputs)
             compiled_layers[layer.name] = model
-            retval = KerasModel(model_inputs, model, name="{}_inference".format(saved_model.name))
+            retval = KerasModel(model_inputs, model, name=f"{saved_model.name}_inference")
         logger.debug("Compiled inference model '%s': %s", retval.name, retval)
         return retval
 

plugins/train/model/dfaker.py

@@ -4,11 +4,18 @@
 import logging
 import sys
 
+from lib.model.nn_blocks import Conv2DOutput, UpscaleBlock, ResidualBlock
+from lib.utils import get_backend
+from .original import Model as OriginalModel, KerasModel
+
+if get_backend() == "amd":
     from keras.initializers import RandomNormal
     from keras.layers import Input, LeakyReLU
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.initializers import RandomNormal  # noqa pylint:disable=import-error,no-name-in-module
+    from tensorflow.keras.layers import Input, LeakyReLU  # noqa pylint:disable=import-error,no-name-in-module
 
-from lib.model.nn_blocks import Conv2DOutput, UpscaleBlock, ResidualBlock
-from .original import Model as OriginalModel, KerasModel
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
@@ -44,7 +51,7 @@ class Model(OriginalModel):
         var_x = LeakyReLU(alpha=0.2)(var_x)
         var_x = ResidualBlock(128, kernel_initializer=self.kernel_initializer)(var_x)
         var_x = UpscaleBlock(64, activation="leakyrelu")(var_x)
-        var_x = Conv2DOutput(3, 5, name="face_out_{}".format(side))(var_x)
+        var_x = Conv2DOutput(3, 5, name=f"face_out_{side}")(var_x)
         outputs = [var_x]
 
         if self.config.get("learn_mask", False):
@@ -55,6 +62,6 @@ class Model(OriginalModel):
             var_y = UpscaleBlock(256, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(128, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(64, activation="leakyrelu")(var_y)
-            var_y = Conv2DOutput(1, 5, name="mask_out_{}".format(side))(var_y)
+            var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y)
             outputs.append(var_y)
-        return KerasModel([input_], outputs=outputs, name="decoder_{}".format(side))
+        return KerasModel([input_], outputs=outputs, name=f"decoder_{side}")

plugins/train/model/dfl_h128.py

@@ -3,11 +3,16 @@
     Based on https://github.com/iperov/DeepFaceLab
 """
 
-from keras.layers import Dense, Flatten, Input, Reshape
-
 from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock
+from lib.utils import get_backend
 from .original import Model as OriginalModel, KerasModel
 
+if get_backend() == "amd":
+    from keras.layers import Dense, Flatten, Input, Reshape
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import Dense, Flatten, Input, Reshape  # noqa pylint:disable=import-error,no-name-in-module
+
 
 class Model(OriginalModel):
     """ H128 Model from DFL """
@@ -36,7 +41,7 @@ class Model(OriginalModel):
         var_x = UpscaleBlock(self.encoder_dim, activation="leakyrelu")(var_x)
         var_x = UpscaleBlock(self.encoder_dim // 2, activation="leakyrelu")(var_x)
         var_x = UpscaleBlock(self.encoder_dim // 4, activation="leakyrelu")(var_x)
-        var_x = Conv2DOutput(3, 5, name="face_out_{}".format(side))(var_x)
+        var_x = Conv2DOutput(3, 5, name=f"face_out_{side}")(var_x)
         outputs = [var_x]
 
         if self.config.get("learn_mask", False):
@@ -44,6 +49,6 @@ class Model(OriginalModel):
             var_y = UpscaleBlock(self.encoder_dim, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(self.encoder_dim // 2, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(self.encoder_dim // 4, activation="leakyrelu")(var_y)
-            var_y = Conv2DOutput(1, 5, name="mask_out_{}".format(side))(var_y)
+            var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y)
             outputs.append(var_y)
-        return KerasModel(input_, outputs=outputs, name="decoder_{}".format(side))
+        return KerasModel(input_, outputs=outputs, name=f"decoder_{side}")

plugins/train/model/dfl_sae.py

@@ -5,12 +5,17 @@
 
 import numpy as np
 
-from keras.layers import Concatenate, Dense, Flatten, Input, LeakyReLU, Reshape
-
 from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock
+from lib.utils import get_backend
 
 from ._base import ModelBase, KerasModel, logger
 
+if get_backend() == "amd":
+    from keras.layers import Concatenate, Dense, Flatten, Input, LeakyReLU, Reshape
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import Concatenate, Dense, Flatten, Input, LeakyReLU, Reshape  # noqa pylint:disable=import-error,no-name-in-module
+
 
 class Model(ModelBase):
     """ SAE Model from DFL """
@@ -50,7 +55,7 @@ class Model(ModelBase):
 
     def build_model(self, inputs):
         """ Build the DFL-SAE Model """
-        encoder = getattr(self, "encoder_{}".format(self.architecture))()
+        encoder = getattr(self, f"encoder_{self.architecture}")()
         enc_output_shape = encoder.output_shape[1:]
         encoder_a = encoder(inputs[0])
         encoder_b = encoder(inputs[1])
@@ -108,7 +113,7 @@ class Model(ModelBase):
         var_x = Dense(lowest_dense_res * lowest_dense_res * self.ae_dims * 2)(var_x)
         var_x = Reshape((lowest_dense_res, lowest_dense_res, self.ae_dims * 2))(var_x)
         var_x = UpscaleBlock(self.ae_dims * 2, activation="leakyrelu")(var_x)
-        return KerasModel(input_, var_x, name="intermediate_{}".format(side))
+        return KerasModel(input_, var_x, name=f"intermediate_{side}")
 
     def decoder(self, side, input_shape):
         """ DFL SAE Decoder Network"""
@@ -123,38 +128,38 @@ class Model(ModelBase):
         var_x1 = ResidualBlock(dims * 8)(var_x1)
         var_x1 = ResidualBlock(dims * 8)(var_x1)
         if self.multiscale_count >= 3:
-            outputs.append(Conv2DOutput(3, 5, name="face_out_32_{}".format(side))(var_x1))
+            outputs.append(Conv2DOutput(3, 5, name=f"face_out_32_{side}")(var_x1))
 
         var_x2 = UpscaleBlock(dims * 4, activation=None)(var_x1)
         var_x2 = LeakyReLU(alpha=0.2)(var_x2)
         var_x2 = ResidualBlock(dims * 4)(var_x2)
         var_x2 = ResidualBlock(dims * 4)(var_x2)
         if self.multiscale_count >= 2:
-            outputs.append(Conv2DOutput(3, 5, name="face_out_64_{}".format(side))(var_x2))
+            outputs.append(Conv2DOutput(3, 5, name=f"face_out_64_{side}")(var_x2))
 
         var_x3 = UpscaleBlock(dims * 2, activation=None)(var_x2)
         var_x3 = LeakyReLU(alpha=0.2)(var_x3)
         var_x3 = ResidualBlock(dims * 2)(var_x3)
         var_x3 = ResidualBlock(dims * 2)(var_x3)
 
-        outputs.append(Conv2DOutput(3, 5, name="face_out_128_{}".format(side))(var_x3))
+        outputs.append(Conv2DOutput(3, 5, name=f"face_out_128_{side}")(var_x3))
 
         if self.use_mask:
             var_y = input_
             var_y = UpscaleBlock(self.decoder_dim * 8, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(self.decoder_dim * 4, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(self.decoder_dim * 2, activation="leakyrelu")(var_y)
-            var_y = Conv2DOutput(1, 5, name="mask_out_{}".format(side))(var_y)
+            var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y)
             outputs.append(var_y)
-        return KerasModel(input_, outputs=outputs, name="decoder_{}".format(side))
+        return KerasModel(input_, outputs=outputs, name=f"decoder_{side}")
 
     def _legacy_mapping(self):
         """ The mapping of legacy separate model names to single model names """
-        mappings = dict(df={"{}_encoder.h5".format(self.name): "encoder_df",
-                            "{}_decoder_A.h5".format(self.name): "decoder_a",
-                            "{}_decoder_B.h5".format(self.name): "decoder_b"},
-                        liae={"{}_encoder.h5".format(self.name): "encoder_liae",
-                              "{}_intermediate_B.h5".format(self.name): "intermediate_both",
-                              "{}_intermediate.h5".format(self.name): "intermediate_b",
-                              "{}_decoder.h5".format(self.name): "decoder_both"})
+        mappings = dict(df={f"{self.name}_encoder.h5": "encoder_df",
+                            f"{self.name}_decoder_A.h5": "decoder_a",
+                            f"{self.name}_decoder_B.h5": "decoder_b"},
+                        liae={f"{self.name}_encoder.h5": "encoder_liae",
+                              f"{self.name}_intermediate_B.h5": "intermediate_both",
+                              f"{self.name}_intermediate.h5": "intermediate_b",
+                              f"{self.name}_decoder.h5": "decoder_both"})
         return mappings[self.config["architecture"]]

plugins/train/model/dlight.py

@@ -8,15 +8,22 @@
     DeepHomage for lots of testing
     """
 
-from keras.layers import (AveragePooling2D, BatchNormalization, Concatenate, Dense, Dropout,
-                          Flatten, Input, Reshape, LeakyReLU, UpSampling2D)
-
 from lib.model.nn_blocks import (Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock,
                                  Upscale2xBlock)
-from lib.utils import FaceswapError
+from lib.utils import FaceswapError, get_backend
 
 from ._base import ModelBase, KerasModel, logger
 
+if get_backend() == "amd":
+    from keras.layers import (
+        AveragePooling2D, BatchNormalization, Concatenate, Dense, Dropout, Flatten, Input, Reshape,
+        LeakyReLU, UpSampling2D)
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import (  # pylint:disable=import-error,no-name-in-module
+        AveragePooling2D, BatchNormalization, Concatenate, Dense, Dropout, Flatten, Input, Reshape,
+        LeakyReLU, UpSampling2D)
+
 
 class Model(ModelBase):
     """ DLight Autoencoder Model """
@@ -218,6 +225,6 @@ class Model(ModelBase):
     def _legacy_mapping(self):
         """ The mapping of legacy separate model names to single model names """
         decoder_b = "decoder_b" if self.details > 0 else "decoder_b_fast"
-        return {"{}_encoder.h5".format(self.name): "encoder",
-                "{}_decoder_A.h5".format(self.name): "decoder_a",
-                "{}_decoder_B.h5".format(self.name): decoder_b}
+        return {f"{self.name}_encoder.h5": "encoder",
+                f"{self.name}_decoder_A.h5": "decoder_a",
+                f"{self.name}_decoder_B.h5": decoder_b}

plugins/train/model/iae.py

@@ -1,10 +1,17 @@
 #!/usr/bin/env python3
 """ Improved autoencoder for faceswap """
 
+from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock
+from lib.utils import get_backend
+
+from ._base import ModelBase, KerasModel
+
+if get_backend() == "amd":
     from keras.layers import Concatenate, Dense, Flatten, Input, Reshape
 
-from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock
-from ._base import ModelBase, KerasModel
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import Concatenate, Dense, Flatten, Input, Reshape  # noqa pylint:disable=import-error,no-name-in-module
 
 
 class Model(ModelBase):
@@ -48,7 +55,7 @@ class Model(ModelBase):
         var_x = Dense(self.encoder_dim)(input_)
         var_x = Dense(4 * 4 * int(self.encoder_dim/2))(var_x)
         var_x = Reshape((4, 4, int(self.encoder_dim/2)))(var_x)
-        return KerasModel(input_, var_x, name="inter_{}".format(side))
+        return KerasModel(input_, var_x, name=f"inter_{side}")
 
     def decoder(self):
         """ Decoder Network """
@@ -73,8 +80,8 @@ class Model(ModelBase):
 
     def _legacy_mapping(self):
         """ The mapping of legacy separate model names to single model names """
-        return {"{}_encoder.h5".format(self.name): "encoder",
-                "{}_intermediate_A.h5".format(self.name): "inter_a",
-                "{}_intermediate_B.h5".format(self.name): "inter_b",
-                "{}_inter.h5".format(self.name): "inter_both",
-                "{}_decoder.h5".format(self.name): "decoder"}
+        return {f"{self.name}_encoder.h5": "encoder",
+                f"{self.name}_intermediate_A.h5": "inter_a",
+                f"{self.name}_intermediate_B.h5": "inter_b",
+                f"{self.name}_inter.h5": "inter_both",
+                f"{self.name}_decoder.h5": "decoder"}

plugins/train/model/lightweight.py

@@ -4,10 +4,8 @@
     Based on the original https://www.reddit.com/r/deepfakes/
     code sample + contributions """
 
-from keras.layers import Dense, Flatten, Input, Reshape
-
 from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock
-from .original import Model as OriginalModel, KerasModel
+from .original import Model as OriginalModel, KerasModel, Dense, Flatten, Input, Reshape
 
 
 class Model(OriginalModel):
@@ -36,7 +34,7 @@ class Model(OriginalModel):
         var_x = UpscaleBlock(512, activation="leakyrelu")(var_x)
         var_x = UpscaleBlock(256, activation="leakyrelu")(var_x)
         var_x = UpscaleBlock(128, activation="leakyrelu")(var_x)
-        var_x = Conv2DOutput(3, 5, activation="sigmoid", name="face_out_{}".format(side))(var_x)
+        var_x = Conv2DOutput(3, 5, activation="sigmoid", name=f"face_out_{side}")(var_x)
         outputs = [var_x]
 
         if self.config.get("learn_mask", False):
@@ -46,6 +44,6 @@ class Model(OriginalModel):
             var_y = UpscaleBlock(128, activation="leakyrelu")(var_y)
             var_y = Conv2DOutput(1, 5,
                                  activation="sigmoid",
-                                 name="mask_out_{}".format(side))(var_y)
+                                 name=f"mask_out_{side}")(var_y)
             outputs.append(var_y)
-        return KerasModel(input_, outputs=outputs, name="decoder_{}".format(side))
+        return KerasModel(input_, outputs=outputs, name=f"decoder_{side}")

plugins/train/model/original.py

@@ -5,11 +5,17 @@ Based on the original https://www.reddit.com/r/deepfakes/ code sample + contribu
 This model is heavily documented as it acts as a template that other model plugins can be developed
 from.
 """
-from keras.layers import Dense, Flatten, Reshape, Input
 
 from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, UpscaleBlock
+from lib.utils import get_backend
 from ._base import KerasModel, ModelBase
 
+if get_backend() == "amd":
+    from keras.layers import Dense, Flatten, Reshape, Input
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.layers import Dense, Flatten, Reshape, Input  # noqa pylint:disable=import-error,no-name-in-module
+
 
 class Model(ModelBase):
     """ Original Faceswap Model.
@@ -144,7 +150,7 @@ class Model(ModelBase):
         var_x = UpscaleBlock(256, activation="leakyrelu")(var_x)
         var_x = UpscaleBlock(128, activation="leakyrelu")(var_x)
         var_x = UpscaleBlock(64, activation="leakyrelu")(var_x)
-        var_x = Conv2DOutput(3, 5, name="face_out_{}".format(side))(var_x)
+        var_x = Conv2DOutput(3, 5, name=f"face_out_{side}")(var_x)
         outputs = [var_x]
 
         if self.learn_mask:
@@ -152,12 +158,12 @@ class Model(ModelBase):
             var_y = UpscaleBlock(256, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(128, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(64, activation="leakyrelu")(var_y)
-            var_y = Conv2DOutput(1, 5, name="mask_out_{}".format(side))(var_y)
+            var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y)
             outputs.append(var_y)
-        return KerasModel(input_, outputs=outputs, name="decoder_{}".format(side))
+        return KerasModel(input_, outputs=outputs, name=f"decoder_{side}")
 
     def _legacy_mapping(self):
         """ The mapping of legacy separate model names to single model names """
-        return {"{}_encoder.h5".format(self.name): "encoder",
-                "{}_decoder_A.h5".format(self.name): "decoder_a",
-                "{}_decoder_B.h5".format(self.name): "decoder_b"}
+        return {f"{self.name}_encoder.h5": "encoder",
+                f"{self.name}_decoder_A.h5": "decoder_a",
+                f"{self.name}_decoder_B.h5": "decoder_b"}

plugins/train/model/phaze_a.py

@@ -2,15 +2,6 @@
 """ Phaze-A Model by TorzDF with thanks to BirbFakes and the myriad of testers. """
 
 import numpy as np
-import tensorflow as tf
-
-import keras.backend as K
-from keras import applications as kapp
-from keras.layers import (
-    Add, BatchNormalization, Concatenate, Dense, Dropout, Flatten, GaussianNoise,
-    GlobalAveragePooling2D, GlobalMaxPooling2D, Input, LeakyReLU, Reshape, UpSampling2D,
-    Conv2D as KConv2D)
-from keras.models import clone_model
 
 from lib.model.nn_blocks import (
     Conv2D, Conv2DBlock, Conv2DOutput, ResidualBlock, UpscaleBlock, Upscale2xBlock,
@@ -18,11 +9,26 @@ from lib.model.nn_blocks import (
 from lib.model.normalization import (
     AdaInstanceNormalization, GroupNormalization, InstanceNormalization, LayerNormalization,
     RMSNormalization)
-
-from lib.utils import get_backend, FaceswapError
+from lib.utils import get_backend, get_tf_version, FaceswapError
 
 from ._base import KerasModel, ModelBase, logger, _get_all_sub_models
 
+if get_backend() == "amd":
+    from keras import applications as kapp, backend as K
+    from keras.layers import (
+        Add, BatchNormalization, Concatenate, Dense, Dropout, Flatten, GaussianNoise,
+        GlobalAveragePooling2D, GlobalMaxPooling2D, Input, LeakyReLU, Reshape, UpSampling2D,
+        Conv2D as KConv2D)
+    from keras.models import clone_model
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import applications as kapp, backend as K  # pylint:disable=import-error
+    from tensorflow.keras.layers import (  # pylint:disable=import-error,no-name-in-module
+        Add, BatchNormalization, Concatenate, Dense, Dropout, Flatten, GaussianNoise,
+        GlobalAveragePooling2D, GlobalMaxPooling2D, Input, LeakyReLU, Reshape, UpSampling2D,
+        Conv2D as KConv2D)
+    from tensorflow.keras.models import clone_model  # noqa pylint:disable=import-error,no-name-in-module
+
 
 _MODEL_MAPPING = dict(
     densenet121=dict(
@@ -47,6 +53,20 @@ _MODEL_MAPPING = dict(
         keras_name="EfficientNetB6", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=528),
     efficientnet_b7=dict(
         keras_name="EfficientNetB7", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=600),
+    efficientnet_v2_b0=dict(
+        keras_name="EfficientNetV2B0", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=224),
+    efficientnet_v2_b1=dict(
+        keras_name="EfficientNetV2B1", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=240),
+    efficientnet_v2_b2=dict(
+        keras_name="EfficientNetV2B2", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=260),
+    efficientnet_v2_b3=dict(
+        keras_name="EfficientNetV2B3", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=300),
+    efficientnet_v2_s=dict(
+        keras_name="EfficientNetV2S", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=384),
+    efficientnet_v2_m=dict(
+        keras_name="EfficientNetV2M", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=480),
+    efficientnet_v2_l=dict(
+        keras_name="EfficientNetV2L", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=480),
     inception_resnet_v2=dict(
         keras_name="InceptionResNetV2", scaling=(-1, 1), min_size=75, default_size=299),
     inception_v3=dict(
@@ -55,6 +75,10 @@ _MODEL_MAPPING = dict(
         keras_name="MobileNet", scaling=(-1, 1), default_size=224),
     mobilenet_v2=dict(
         keras_name="MobileNetV2", scaling=(-1, 1), default_size=224),
+    mobilenet_v3_large=dict(
+        keras_name="MobileNetV3Large", no_amd=True, tf_min=2.4, scaling=(-1, 1), default_size=224),
+    mobilenet_v3_small=dict(
+        keras_name="MobileNetV3Small", no_amd=True, tf_min=2.4, scaling=(-1, 1), default_size=224),
     nasnet_large=dict(
         keras_name="NASNetLarge", scaling=(-1, 1), default_size=331, enforce_for_weights=True),
     nasnet_mobile=dict(
@@ -182,17 +206,20 @@ class Model(ModelBase):
         list
             The selected layers for weight freezing
         """
+        arch = self.config["enc_architecture"]
         layers = self.config["freeze_layers"]
-        keras_name = _MODEL_MAPPING[self.config["enc_architecture"]].get("keras_name")
+        # EfficientNetV2 is inconsistent with other model's naming conventions
+        keras_name = _MODEL_MAPPING[arch].get("keras_name").replace("EfficientNetV2",
+                                                                    "EfficientNetV2-")
 
         if "keras_encoder" not in self.config["freeze_layers"]:
             retval = layers
         elif keras_name:
             retval = [layer.replace("keras_encoder", keras_name.lower()) for layer in layers]
-            logger.debug("Substituting 'keras_encoder' for '%s'", self.config["enc_architecture"])
+            logger.debug("Substituting 'keras_encoder' for '%s'", arch)
         else:
             retval = [layer for layer in layers if layer != "keras_encoder"]
-            logger.debug("Removing 'keras_encoder' for '%s'", self.config["enc_architecture"])
+            logger.debug("Removing 'keras_encoder' for '%s'", arch)
         return retval
 
     def _get_input_shape(self):
@@ -201,16 +228,32 @@ class Model(ModelBase):
         Input shape is calculated from the selected Encoder's input size, scaled to the user
         selected Input Scaling, rounded down to the nearest 16 pixels.
 
+        Notes
+        -----
+        Some models (NasNet) require the input size to be of a certain dimension if loading
+        imagenet weights. In these instances resize inputs and raise warning message
+
         Returns
         -------
         tuple
             The shape tuple for the input size to the Phaze-A model
         """
-        size = _MODEL_MAPPING[self.config["enc_architecture"]]["default_size"]
-        min_size = _MODEL_MAPPING[self.config["enc_architecture"]].get("min_size", 32)
+        arch = self.config["enc_architecture"]
+        enforce_size = _MODEL_MAPPING[arch].get("enforce_for_weights", False)
+        default_size = _MODEL_MAPPING[arch]["default_size"]
         scaling = self.config["enc_scaling"] / 100
-        size = int(max(min_size, min(size, ((size * scaling) // 16) * 16)))
+
+        min_size = _MODEL_MAPPING[arch].get("min_size", 32)
+        size = int(max(min_size, min(default_size, ((default_size * scaling) // 16) * 16)))
+
+        if self.config["enc_load_weights"] and enforce_size and scaling != 1.0:
+            logger.warning("%s requires input size to be %spx when loading imagenet weights. "
+                           "Adjusting input size from %spx to %spx",
+                           arch, default_size, size, default_size)
+            retval = (default_size, default_size, 3)
+        else:
             retval = (size, size, 3)
+
         logger.debug("Encoder input set to: %s", retval)
         return retval
 
@@ -227,11 +270,11 @@ class Model(ModelBase):
                                 f"one of {list(_MODEL_MAPPING.keys())}.")
 
         if get_backend() == "amd" and model.get("no_amd"):
-            valid = [x for x in _MODEL_MAPPING if not _MODEL_MAPPING[x].get('no_amd')]
+            valid = [k for k, v in _MODEL_MAPPING.items() if not v.get('no_amd')]
             raise FaceswapError(f"'{arch}' is not compatible with the AMD backend. Choose one of "
                                 f"{valid}.")
 
-        tf_ver = float(".".join(tf.__version__.split(".")[:2]))  # pylint:disable=no-member
+        tf_ver = get_tf_version()
         tf_min = model.get("tf_min", 2.0)
         if get_backend() != "amd" and tf_ver < tf_min:
             raise FaceswapError(f"{arch}' is not compatible with your version of Tensorflow. The "
@@ -542,32 +585,21 @@ class Encoder():  # pylint:disable=too-few-public-methods
         return dict(mobilenet=dict(alpha=self._config["mobilenet_width"],
                                    depth_multiplier=self._config["mobilenet_depth"],
                                    dropout=self._config["mobilenet_dropout"]),
-                    mobilenet_v2=dict(alpha=self._config["mobilenet_width"]))
+                    mobilenet_v2=dict(alpha=self._config["mobilenet_width"]),
+                    mobilenet_v3=dict(alpha=self._config["mobilenet_width"],
+                                      minimalist=self._config["mobilenet_minimalistic"],
+                                      include_preprocessing=False))
 
     @property
     def _selected_model(self):
         """ dict: The selected encoder model options dictionary """
         arch = self._config["enc_architecture"]
         model = _MODEL_MAPPING.get(arch)
-        model["kwargs"] = self._model_kwargs.get(arch, dict())
+        model["kwargs"] = self._model_kwargs.get(arch, {})
+        if arch.startswith("efficientnet_v2"):
+            model["kwargs"]["include_preprocessing"] = False
         return model
 
-    @property
-    def _model_input_shape(self):
-        """ tuple: The required input shape for the encoder model.
-
-        Notes
-        -----
-        NasNet does not allow custom input sizes when loading pre-trained weights, so we need to
-        resize the input for this model
-        """
-        default_size = self._selected_model.get("default_size")
-        if self._config["enc_load_weights"] and self._selected_model.get("enforce_for_weights"):
-            retval = (default_size, default_size, 3)
-        else:
-            retval = self._input_shape
-        return retval
-
     def __call__(self):
         """ Create the Phaze-A Encoder Model.
 
@@ -576,12 +608,9 @@ class Encoder():  # pylint:disable=too-few-public-methods
         :class:`keras.models.Model`
             The selected Encoder Model
         """
-        input_ = Input(shape=self._model_input_shape)
+        input_ = Input(shape=self._input_shape)
         var_x = input_
 
-        if self._input_shape != self._model_input_shape:
-            var_x = self._resize_inputs(var_x)
-
         scaling = self._selected_model.get("scaling")
         if scaling:
             #  Some models expect different scaling.
@@ -604,28 +633,6 @@ class Encoder():  # pylint:disable=too-few-public-methods
 
         return KerasModel(input_, var_x, name="encoder")
 
-    def _resize_inputs(self, inputs):
-        """ Some models (specifically NasNet) need a specific input size when loading trained
-        weights. This is slightly hacky, but arbitrarily resize the input for these instances.
-
-        Parameters
-        ----------
-        inputs: tensor
-            The input tensor to be resized
-
-        Returns
-        -------
-        tensor
-            The resized input tensor
-        """
-        input_size = self._input_shape[0]
-        new_size = self._model_input_shape[0]
-        logger.debug("Resizing input for encoder: '%s' from %s to %s due to trained weights usage",
-                     self._config["enc_architecture"], input_size, new_size)
-        scale = new_size / input_size
-        interp = "bilinear" if scale > 1 else "nearest"
-        return K.resize_images(size=scale, interpolation=interp)(inputs)
-
     def _get_encoder_model(self):
         """ Return the model defined by the selected architecture.
 
@@ -641,7 +648,7 @@ class Encoder():  # pylint:disable=too-few-public-methods
         """
         if self._selected_model.get("keras_name"):
             kwargs = self._selected_model["kwargs"]
-            kwargs["input_shape"] = self._model_input_shape
+            kwargs["input_shape"] = self._input_shape
             kwargs["include_top"] = False
             kwargs["weights"] = "imagenet" if self._config["enc_load_weights"] else None
             retval = getattr(kapp, self._selected_model["keras_name"])(**kwargs)
@@ -800,7 +807,7 @@ class FullyConnected():  # pylint:disable=too-few-public-methods
                 if self._config["fc_upsampler"].lower() == "upsample2d":
                     var_x = LeakyReLU(alpha=0.1)(var_x)
 
-        return KerasModel(input_, var_x, name="fc_{}".format(self._side))
+        return KerasModel(input_, var_x, name=f"fc_{self._side}")
 
 
 class GBlock():  # pylint:disable=too-few-public-methods
@@ -1024,4 +1031,4 @@ class Decoder():  # pylint:disable=too-few-public-methods
                                         self._config["dec_output_kernel"],
                                         name="mask_out")(var_y))
 
-        return KerasModel(inputs, outputs=outputs, name="decoder_{}".format(self._side))
+        return KerasModel(inputs, outputs=outputs, name=f"decoder_{self._side}")

plugins/train/model/phaze_a_defaults.py

@@ -52,6 +52,9 @@ _ENCODERS = ["densenet121", "densenet169", "densenet201", "inception_resnet_v2",
 if get_backend() != "amd":
     _ENCODERS.extend(["efficientnet_b0", "efficientnet_b1", "efficientnet_b2", "efficientnet_b3",
                       "efficientnet_b4", "efficientnet_b5", "efficientnet_b6", "efficientnet_b7",
+                      "efficientnet_v2_b0", "efficientnet_v2_b1", "efficientnet_v2_b2",
+                      "efficientnet_v2_b3", "efficientnet_v2_l", "efficientnet_v2_m",
+                      "efficientnet_v2_s", "mobilenet_v3_large", "mobilenet_v3_small",
                       "resnet50_v2", "resnet101", "resnet101_v2", "resnet152", "resnet152_v2"])
 _ENCODERS = sorted(_ENCODERS)
 
@@ -142,6 +145,13 @@ _DEFAULTS = dict(
              "each variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, b4: 380px, b5: 456px, "
              "b6: 528px, b7 600px. Ref: Rethinking Model Scaling for Convolutional Neural "
              "Networks (2020): https://arxiv.org/abs/1905.11946"
+             "\n\tefficientnet_v2: [Tensorflow 2.8+ only] EfficientNetV2 is the follow up to "
+             "efficientnet. It has numerous variants (B0 - B3 and Small, Medium and Large) that "
+             "increases the model width, depth and dimensional space at each step. The minimum "
+             "input resolution is 32px for all variants. The maximum input resolution for each "
+             "variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, s: 384px, m: 480px, l: "
+             "480px. Ref: EfficientNetV2: Smaller Models and Faster Training (2021): "
+             "https://arxiv.org/abs/2104.00298"
              "\n\tfs_original: (32px - 160px). A configurable variant of the original facewap "
              "encoder. ImageNet weights cannot be loaded for this model. Additional parameters "
              "can be configured with the 'fs_enc' options. A version of this encoder is used in "
@@ -157,6 +167,9 @@ _DEFAULTS = dict(
              "\n\tmobilenet_v2: (32px - 224px). Additional MobileNet parameters can be set with "
              "the 'mobilenet' options. Ref: MobileNetV2: Inverted Residuals and Linear "
              "Bottlenecks (2018): https://arxiv.org/abs/1801.04381"
+             "\n\tmobilenet_v3: (32px - 224px). Additional MobileNet parameters can be set with "
+             "the 'mobilenet' options. Ref: Searching for MobileNetV3 (2019): "
+             "https://arxiv.org/pdf/1905.02244.pdf"
              "\n\tnasnet: (32px - 331px (large) or 224px (mobile)). Ref: Learning Transferable "
              "Architectures for Scalable Image Recognition (2017): "
              "https://arxiv.org/abs/1707.07012"
@@ -569,9 +582,10 @@ _DEFAULTS = dict(
              "each layer. Values greater than 1.0 proportionally increase the number of filters "
              "within each layer. 1.0 is the default number of layers used within the paper.\n"
              "NB: This option is ignored for any non-mobilenet encoders.\n"
-             "NB: If loading ImageNet weights, then for mobilenet v1 only values of '0.25', "
-             "'0.5', '0.75' or '1.0 can be selected. For mobilenet v2 only values of '0.35', "
-             "'0.50', '0.75', '1.0', '1.3' or '1.4' can be selected",
+             "NB: If loading ImageNet weights, then for MobilenetV1 only values of '0.25', "
+             "'0.5', '0.75' or '1.0 can be selected. For MobilenetV2 only values of '0.35', "
+             "'0.50', '0.75', '1.0', '1.3' or '1.4' can be selected. For mobilenet_v3 only values "
+             "of '0.75' or '1.0' can be selected",
         datatype=float,
         min_max=(0.1, 2.0),
         rounding=2,
@@ -579,10 +593,10 @@ _DEFAULTS = dict(
         fixed=True),
     mobilenet_depth=dict(
         default=1,
-        info="The depth multiplier for mobilenet v1 encoder. This is the depth multiplier "
+        info="The depth multiplier for MobilenetV1 encoder. This is the depth multiplier "
              "for depthwise convolution (known as the resolution multiplier within the original "
              "paper).\n"
-             "NB: This option is only used for mobilenet v1 and is ignored for all other "
+             "NB: This option is only used for MobilenetV1 and is ignored for all other "
              "encoders.\n"
              "NB: If loading ImageNet weights, this must be set to 1.",
         datatype=int,
@@ -592,13 +606,25 @@ _DEFAULTS = dict(
         fixed=True),
     mobilenet_dropout=dict(
         default=0.001,
-        info="The dropout rate for for mobilenet v1 encoder.\n"
-             "NB: This option is only used for mobilenet v1 and is ignored for all other "
-             "encoders.\n"
-             "NB: If loading ImageNet weights, this must be set to 1.0.",
+        info="The dropout rate for MobilenetV1 encoder.\n"
+             "NB: This option is only used for MobilenetV1 and is ignored for all other "
+             "encoders.",
         datatype=float,
-        min_max=(0.1, 2.0),
-        rounding=2,
+        min_max=(0.001, 2.0),
+        rounding=3,
+        group="mobilenet encoder configuration",
+        fixed=True),
+    mobilenet_minimalistic=dict(
+        default=False,
+        info="Use a minimilist version of MobilenetV3.\n"
+             "In addition to large and small models MobilenetV3 also contains so-called "
+             "minimalistic models, these models have the same per-layer dimensions characteristic "
+             "as MobilenetV3 however, they don't utilize any of the advanced blocks "
+             "(squeeze-and-excite units, hard-swish, and 5x5 convolutions). While these models "
+             "are less efficient on CPU, they are much more performant on GPU/DSP.\n"
+             "NB: This option is only used for MobilenetV3 and is ignored for all other "
+             "encoders.\n",
+        datatype=bool,
         group="mobilenet encoder configuration",
         fixed=True),
     )

plugins/train/model/realface.py

@@ -9,12 +9,17 @@
     """
 import sys
 
+from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock
+from lib.utils import get_backend
+from ._base import ModelBase, KerasModel, logger
+
+if get_backend() == "amd":
     from keras.initializers import RandomNormal
     from keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape
-
-
-from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock
-from ._base import ModelBase, KerasModel, logger
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.initializers import RandomNormal  # noqa pylint:disable=import-error,no-name-in-module
+    from tensorflow.keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape  # noqa pylint:disable=import-error,no-name-in-module
 
 
 class Model(ModelBase):
@@ -183,6 +188,6 @@ class Model(ModelBase):
 
     def _legacy_mapping(self):
         """ The mapping of legacy separate model names to single model names """
-        return {"{}_encoder.h5".format(self.name): "encoder",
-                "{}_decoder_A.h5".format(self.name): "decoder_a",
-                "{}_decoder_B.h5".format(self.name): "decoder_b"}
+        return {f"{self.name}_encoder.h5": "encoder",
+                f"{self.name}_decoder_A.h5": "decoder_a",
+                f"{self.name}_decoder_B.h5": "decoder_b"}

plugins/train/model/unbalanced.py

@@ -3,11 +3,17 @@
     Based on the original https://www.reddit.com/r/deepfakes/
         code sample + contributions """
 
+from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock
+from lib.utils import get_backend
+from ._base import ModelBase, KerasModel
+
+if get_backend() == "amd":
     from keras.initializers import RandomNormal
     from keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape, SpatialDropout2D
-
-from lib.model.nn_blocks import Conv2DOutput, Conv2DBlock, ResidualBlock, UpscaleBlock
-from ._base import ModelBase, KerasModel
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.initializers import RandomNormal  # noqa pylint:disable=import-error,no-name-in-module
+    from tensorflow.keras.layers import Dense, Flatten, Input, LeakyReLU, Reshape, SpatialDropout2D  # noqa pylint:disable=import-error,no-name-in-module
 
 
 class Model(ModelBase):
@@ -135,6 +141,6 @@ class Model(ModelBase):
 
     def _legacy_mapping(self):
         """ The mapping of legacy separate model names to single model names """
-        return {"{}_encoder.h5".format(self.name): "encoder",
-                "{}_decoder_A.h5".format(self.name): "decoder_a",
-                "{}_decoder_B.h5".format(self.name): "decoder_b"}
+        return {f"{self.name}_encoder.h5": "encoder",
+                f"{self.name}_decoder_A.h5": "decoder_a",
+                f"{self.name}_decoder_B.h5": "decoder_b"}

plugins/train/model/villain.py

@@ -3,14 +3,21 @@
     Based on the original https://www.reddit.com/r/deepfakes/ code sample + contributions
     Adapted from a model by VillainGuy (https://github.com/VillainGuy) """
 
-from keras.initializers import RandomNormal
-from keras.layers import add, Dense, Flatten, Input, LeakyReLU, Reshape
-
 from lib.model.layers import PixelShuffler
 from lib.model.nn_blocks import (Conv2DOutput, Conv2DBlock, ResidualBlock, SeparableConv2DBlock,
                                  UpscaleBlock)
+from lib.utils import get_backend
+
 from .original import Model as OriginalModel, KerasModel
 
+if get_backend() == "amd":
+    from keras.initializers import RandomNormal
+    from keras.layers import add, Dense, Flatten, Input, LeakyReLU, Reshape
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras.initializers import RandomNormal  # noqa pylint:disable=import-error,no-name-in-module
+    from tensorflow.keras.layers import add, Dense, Flatten, Input, LeakyReLU, Reshape  # noqa pylint:disable=import-error,no-name-in-module
+
 
 class Model(OriginalModel):
     """ Villain Faceswap Model """
@@ -72,7 +79,7 @@ class Model(OriginalModel):
         var_x = UpscaleBlock(self.input_shape[0], activation=None, **kwargs)(var_x)
         var_x = LeakyReLU(alpha=0.2)(var_x)
         var_x = ResidualBlock(self.input_shape[0], **kwargs)(var_x)
-        var_x = Conv2DOutput(3, 5, name="face_out_{}".format(side))(var_x)
+        var_x = Conv2DOutput(3, 5, name=f"face_out_{side}")(var_x)
         outputs = [var_x]
 
         if self.config.get("learn_mask", False):
@@ -80,6 +87,6 @@ class Model(OriginalModel):
             var_y = UpscaleBlock(512, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(256, activation="leakyrelu")(var_y)
             var_y = UpscaleBlock(self.input_shape[0], activation="leakyrelu")(var_y)
-            var_y = Conv2DOutput(1, 5, name="mask_out_{}".format(side))(var_y)
+            var_y = Conv2DOutput(1, 5, name=f"mask_out_{side}")(var_y)
             outputs.append(var_y)
-        return KerasModel(input_, outputs=outputs, name="decoder_{}".format(side))
+        return KerasModel(input_, outputs=outputs, name=f"decoder_{side}")

plugins/train/trainer/_base.py

@@ -16,10 +16,11 @@ import cv2
 import numpy as np
 
 import tensorflow as tf
-from tensorflow.python.framework import errors_impl as tf_errors
+from tensorflow.python.framework import (  # pylint:disable=no-name-in-module
+    errors_impl as tf_errors)
 
 from lib.training import TrainingDataGenerator
-from lib.utils import FaceswapError, get_backend, get_folder, get_image_paths
+from lib.utils import FaceswapError, get_backend, get_folder, get_image_paths, get_tf_version
 from plugins.train._config import Config
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
@@ -129,8 +130,8 @@ class TrainerBase():
 
         logger.debug("Setting up TensorBoard Logging")
         log_dir = os.path.join(str(self._model.model_dir),
-                               "{}_logs".format(self._model.name),
-                               "session_{}".format(self._model.state.session_id))
+                               f"{self._model.name}_logs",
+                               f"session_{self._model.state.session_id}")
         tensorboard = tf.keras.callbacks.TensorBoard(log_dir=log_dir,
                                                      histogram_freq=0,  # Must be 0 or hangs
                                                      write_graph=get_backend() != "amd",
@@ -251,7 +252,16 @@ class TrainerBase():
         logger.trace("Updating TensorBoard log")
         logs = {log[0]: log[1]
                 for log in zip(self._model.state.loss_names, loss)}
+
         self._tensorboard.on_train_batch_end(self._model.iterations, logs=logs)
+        if get_tf_version() == 2.8:
+            # Bug in TF 2.8 where batch recording got deleted.
+            # ref: https://github.com/keras-team/keras/issues/16173
+            for name, value in logs.items():
+                tf.summary.scalar(
+                    "batch_" + name,
+                    value,
+                    step=self._model._model._train_counter)  # pylint:disable=protected-access
 
     def _collate_and_store_loss(self, loss):
         """ Collate the loss into totals for each side.
@@ -297,11 +307,11 @@ class TrainerBase():
             The loss for each side. List should contain 2 ``floats`` side "a" in position 0 and
             side "b" in position `.
          """
-        output = ", ".join(["Loss {}: {:.5f}".format(side, side_loss)
+        output = ", ".join([f"Loss {side}: {side_loss:.5f}"
                             for side, side_loss in zip(("A", "B"), loss)])
         timestamp = time.strftime("%H:%M:%S")
-        output = "[{}] [#{:05d}] {}".format(timestamp, self._model.iterations, output)
-        print("\r{}".format(output), end="")
+        output = f"[{timestamp}] [#{self._model.iterations:05d}] {output}"
+        print(f"\r{output}", end="")
 
     def clear_tensorboard(self):
         """ Stop Tensorboard logging.
@@ -335,14 +345,14 @@ class _Feeder():
         self._model = model
         self._images = images
         self._config = config
-        self._target = dict()
-        self._samples = dict()
-        self._masks = dict()
+        self._target = {}
+        self._samples = {}
+        self._masks = {}
 
         self._feeds = {side: self._load_generator(idx).minibatch_ab(images[side], batch_size, side)
                        for idx, side in enumerate(("a", "b"))}
 
-        self._display_feeds = dict(preview=self._set_preview_feed(), timelapse=dict())
+        self._display_feeds = dict(preview=self._set_preview_feed(), timelapse={})
         logger.debug("Initialized %s:", self.__class__.__name__)
 
     def _load_generator(self, output_index):
@@ -385,7 +395,7 @@ class _Feeder():
             The side ("a" or "b") as key, :class:`~lib.training_data.TrainingDataGenerator` as
             value.
         """
-        retval = dict()
+        retval = {}
         for idx, side in enumerate(("a", "b")):
             logger.debug("Setting preview feed: (side: '%s')", side)
             preview_images = self._config.get("preview_images", 14)
@@ -484,9 +494,9 @@ class _Feeder():
             should not be generated, in which case currently stored previews should be deleted.
         """
         if not do_preview:
-            self._samples = dict()
-            self._target = dict()
-            self._masks = dict()
+            self._samples = {}
+            self._target = {}
+            self._masks = {}
             return
         logger.debug("Generating preview")
         for side in ("a", "b"):
@@ -523,7 +533,7 @@ class _Feeder():
          """
         num_images = self._config.get("preview_images", 14)
         num_images = min(batch_size, num_images) if batch_size is not None else num_images
-        retval = dict()
+        retval = {}
         for side in ("a", "b"):
             logger.debug("Compiling samples: (side: '%s', samples: %s)", side, num_images)
             side_images = images[side] if images is not None else self._target[side]
@@ -544,9 +554,9 @@ class _Feeder():
             :class:`numpy.ndarrays` for creating a time-lapse frame
         """
         batchsizes = []
-        samples = dict()
-        images = dict()
-        masks = dict()
+        samples = {}
+        images = {}
+        masks = {}
         for side in ("a", "b"):
             batch = next(self._display_feeds["timelapse"][side])
             batchsizes.append(len(batch["samples"]))
@@ -607,7 +617,7 @@ class _Samples():  # pylint:disable=too-few-public-methods
                      self.__class__.__name__, model, coverage_ratio)
         self._model = model
         self._display_mask = model.config["learn_mask"] or model.config["penalized_mask_loss"]
-        self.images = dict()
+        self.images = {}
         self._coverage_ratio = coverage_ratio
         self._scaling = scaling
         logger.debug("Initialized %s", self.__class__.__name__)
@@ -630,9 +640,9 @@ class _Samples():  # pylint:disable=too-few-public-methods
             A compiled preview image ready for display or saving
         """
         logger.debug("Showing sample")
-        feeds = dict()
-        figures = dict()
-        headers = dict()
+        feeds = {}
+        figures = {}
+        headers = {}
         for idx, side in enumerate(("a", "b")):
             samples = self.images[side]
             faces = samples[1]
@@ -647,8 +657,8 @@ class _Samples():  # pylint:disable=too-few-public-methods
 
         for side, samples in self.images.items():
             other_side = "a" if side == "b" else "b"
-            predictions = [preds["{0}_{0}".format(side)],
-                           preds["{}_{}".format(other_side, side)]]
+            predictions = [preds[f"{side}_{side}"],
+                           preds[f"{other_side}_{side}"]]
             display = self._to_full_frame(side, samples, predictions)
             headers[side] = self._get_headers(side, display[0].shape[1])
             figures[side] = np.stack([display[0], display[1], display[2], ], axis=1)
@@ -716,7 +726,7 @@ class _Samples():  # pylint:disable=too-few-public-methods
             List of :class:`numpy.ndarray` of predictions received from the model
         """
         logger.debug("Getting Predictions")
-        preds = dict()
+        preds = {}
         standard = self._model.model.predict([feed_a, feed_b])
         swapped = self._model.model.predict([feed_b, feed_a])
 
@@ -904,9 +914,9 @@ class _Samples():  # pylint:disable=too-few-public-methods
         total_width = width * 3
         logger.debug("height: %s, total_width: %s", height, total_width)
         font = cv2.FONT_HERSHEY_SIMPLEX
-        texts = ["{} ({})".format(titles[0], side),
-                 "{0} > {0}".format(titles[0]),
-                 "{} > {}".format(titles[0], titles[1])]
+        texts = [f"{titles[0]} ({side})",
+                 f"{titles[0]} > {titles[0]}",
+                 f"{titles[0]} > {titles[1]}"]
         scaling = (width / 144) * 0.45
         text_sizes = [cv2.getTextSize(texts[idx], font, scaling, 1)[0]
                       for idx in range(len(texts))]
@@ -1002,7 +1012,7 @@ class _Timelapse():  # pylint:disable=too-few-public-methods
         logger.debug("Time-lapse output set to '%s'", self._output_file)
 
         # Rewrite paths to pull from the training images so mask and face data can be accessed
-        images = dict()
+        images = {}
         for side, input_ in zip(("a", "b"), (input_a, input_b)):
             training_path = os.path.dirname(self._image_paths[side][0])
             images[side] = [os.path.join(training_path, os.path.basename(pth))

requirements_cpu.txt

@@ -1,2 +1,2 @@
 -r _requirements_base.txt
-tensorflow>=2.2.0,<2.7.0
+tensorflow>=2.2.0,<2.9.0

requirements_nvidia.txt

@@ -1,2 +1,2 @@
 -r _requirements_base.txt
-tensorflow-gpu>=2.2.0,<2.7.0
+tensorflow-gpu>=2.2.0,<2.9.0

scripts/train.py

@@ -360,9 +360,10 @@ class Train():  # pylint:disable=too-few-public-methods
         logger.info("  Starting")
         if self._args.preview:
             logger.info("  Using live preview")
+        if sys.stdout.isatty():
             logger.info("  Press '%s' to save and quit",
                     "Stop" if self._args.redirect_gui or self._args.colab else "ENTER")
-        if not self._args.redirect_gui and not self._args.colab:
+        if not self._args.redirect_gui and not self._args.colab and sys.stdout.isatty():
             logger.info("  Press 'S' to save model weights immediately")
         logger.info("===================================================")
 

setup.py

@@ -19,7 +19,7 @@ INSTALL_FAILED = False
 # Tensorflow builds available from pypi
 TENSORFLOW_REQUIREMENTS = {">=2.2.0,<2.4.0": ["10.1", "7.6"],
                            ">=2.4.0,<2.5.0": ["11.0", "8.0"],
-                           ">=2.5.0,<2.7.0": ["11.2", "8.1"]}
+                           ">=2.5.0,<2.9.0": ["11.2", "8.1"]}
 # Mapping of Python packages to their conda names if different from pip or in non-default channel
 CONDA_MAPPING = {
     # "opencv-python": ("opencv", "conda-forge"),  # Periodic issues with conda-forge opencv
@@ -43,9 +43,9 @@ class Environment():
         self.enable_amd = False
         self.enable_docker = False
         self.enable_cuda = False
-        self.required_packages = list()
-        self.missing_packages = list()
-        self.conda_missing_packages = list()
+        self.required_packages = []
+        self.missing_packages = []
+        self.conda_missing_packages = []
 
         self.process_arguments()
         self.check_permission()
@@ -54,6 +54,7 @@ class Environment():
         self.output_runtime_info()
         self.check_pip()
         self.upgrade_pip()
+        self.set_ld_library_path()
 
         self.installed_packages = self.get_installed_packages()
         self.installed_packages.update(self.get_installed_conda_packages())
@@ -104,7 +105,7 @@ class Environment():
         args = [arg for arg in sys.argv]  # pylint:disable=unnecessary-comprehension
         if self.updater:
             from lib.utils import get_backend  # pylint:disable=import-outside-toplevel
-            args.append("--{}".format(get_backend()))
+            args.append(f"--{get_backend()}")
 
         for arg in args:
             if arg == "--installer":
@@ -124,11 +125,11 @@ class Environment():
             suffix = "cpu.txt"
         req_files = ["_requirements_base.txt", f"requirements_{suffix}"]
         pypath = os.path.dirname(os.path.realpath(__file__))
-        requirements = list()
-        git_requirements = list()
+        requirements = []
+        git_requirements = []
         for req_file in req_files:
             requirements_file = os.path.join(pypath, req_file)
-            with open(requirements_file) as req:
+            with open(requirements_file, encoding="utf8") as req:
                 for package in req.readlines():
                     package = package.strip()
                     # parse_requirements can't handle git dependencies, so extract and then
@@ -157,15 +158,14 @@ class Environment():
         if not self.updater:
             self.output.info("The tool provides tips for installation\n"
                              "and installs required python packages")
-        self.output.info("Setup in %s %s" % (self.os_version[0], self.os_version[1]))
+        self.output.info(f"Setup in {self.os_version[0]} {self.os_version[1]}")
         if not self.updater and not self.os_version[0] in ["Windows", "Linux", "Darwin"]:
-            self.output.error("Your system %s is not supported!" % self.os_version[0])
+            self.output.error(f"Your system {self.os_version[0]} is not supported!")
             sys.exit(1)
 
     def check_python(self):
         """ Check python and virtual environment status """
-        self.output.info("Installed Python: {0} {1}".format(self.py_version[0],
-                                                            self.py_version[1]))
+        self.output.info(f"Installed Python: {self.py_version[0]} {self.py_version[1]}")
         if not (self.py_version[0].split(".")[0] == "3"
                 and self.py_version[0].split(".")[1] in ("7", "8")
                 and self.py_version[1] == "64bit") and not self.updater:
@@ -179,7 +179,7 @@ class Environment():
             self.output.info("Running in Conda")
         if self.is_virtualenv:
             self.output.info("Running in a Virtual Environment")
-        self.output.info("Encoding: {}".format(self.encoding))
+        self.output.info(f"Encoding: {self.encoding}")
 
     def check_pip(self):
         """ Check installed pip version """
@@ -201,16 +201,15 @@ class Environment():
             if not self.is_admin and not self.is_virtualenv:
                 pipexe.append("--user")
             pipexe.append("pip")
-            run(pipexe)
+            run(pipexe, check=True)
         import pip  # pylint:disable=import-outside-toplevel
         pip_version = pip.__version__
-        self.output.info("Installed pip: {}".format(pip_version))
+        self.output.info(f"Installed pip: {pip_version}")
 
     def get_installed_packages(self):
         """ Get currently installed packages """
-        installed_packages = dict()
-        chk = Popen("\"{}\" -m pip freeze".format(sys.executable),
-                    shell=True, stdout=PIPE)
+        installed_packages = {}
+        with Popen(f"\"{sys.executable}\" -m pip freeze", shell=True, stdout=PIPE) as chk:
             installed = chk.communicate()[0].decode(self.encoding).splitlines()
 
         for pkg in installed:
@@ -227,7 +226,7 @@ class Environment():
         chk = os.popen("conda list").read()
         installed = [re.sub(" +", " ", line.strip())
                      for line in chk.splitlines() if not line.startswith("#")]
-        retval = dict()
+        retval = {}
         for pkg in installed:
             item = pkg.split(" ")
             retval[item[0]] = item[1]
@@ -253,7 +252,7 @@ class Environment():
             # that corresponds to the installed Cuda/cuDNN versions
             self.required_packages = [pkg for pkg in self.required_packages
                                       if not pkg.startswith("tensorflow-gpu")]
-            tf_ver = "tensorflow-gpu{}".format(tf_ver)
+            tf_ver = f"tensorflow-gpu{tf_ver}"
             self.required_packages.append(tf_ver)
             return
 
@@ -262,13 +261,12 @@ class Environment():
             "Tensorflow currently has no official prebuild for your CUDA, cuDNN "
             "combination.\nEither install a combination that Tensorflow supports or "
             "build and install your own tensorflow-gpu.\r\n"
-            "CUDA Version: {}\r\n"
-            "cuDNN Version: {}\r\n"
+            f"CUDA Version: {self.cuda_version}\r\n"
+            f"cuDNN Version: {self.cudnn_version}\r\n"
             "Help:\n"
             "Building Tensorflow: https://www.tensorflow.org/install/install_sources\r\n"
             "Tensorflow supported versions: "
-            "https://www.tensorflow.org/install/source#tested_build_configurations".format(
-                self.cuda_version, self.cudnn_version))
+            "https://www.tensorflow.org/install/source#tested_build_configurations")
 
         custom_tf = input("Location of custom tensorflow-gpu wheel (leave "
                           "blank to manually install): ")
@@ -277,9 +275,9 @@ class Environment():
 
         custom_tf = os.path.realpath(os.path.expanduser(custom_tf))
         if not os.path.isfile(custom_tf):
-            self.output.error("{} not found".format(custom_tf))
+            self.output.error(f"{custom_tf} not found")
         elif os.path.splitext(custom_tf)[1] != ".whl":
-            self.output.error("{} is not a valid pip wheel".format(custom_tf))
+            self.output.error(f"{custom_tf} is not a valid pip wheel")
         elif custom_tf:
             self.required_packages.append(custom_tf)
 
@@ -296,9 +294,57 @@ class Environment():
         config = {"backend": backend}
         pypath = os.path.dirname(os.path.realpath(__file__))
         config_file = os.path.join(pypath, "config", ".faceswap")
-        with open(config_file, "w") as cnf:
+        with open(config_file, "w", encoding="utf8") as cnf:
             json.dump(config, cnf)
-        self.output.info("Faceswap config written to: {}".format(config_file))
+        self.output.info(f"Faceswap config written to: {config_file}")
+
+    def set_ld_library_path(self):
+        """ Update the LD_LIBRARY_PATH environment variable when activating a conda environment
+        and revert it when deactivating.
+
+        Notes
+        -----
+        From Tensorflow 2.7, installing Cuda Toolkit from conda-forge and tensorflow from pip
+        causes tensorflow to not be able to locate shared libs and hence not use the GPU.
+        We update the environment variable for all instances using Conda as it shouldn't hurt
+        anything and may help avoid conflicts with globally installed Cuda
+        """
+        if not self.is_conda or not self.enable_cuda:
+            return
+
+        if self.os_version[0] == "Windows":
+            return
+
+        conda_prefix = os.environ["CONDA_PREFIX"]
+        activate_folder = os.path.join(conda_prefix, "etc", "conda", "activate.d")
+        deactivate_folder = os.path.join(conda_prefix, "etc", "conda", "deactivate.d")
+
+        os.makedirs(activate_folder, exist_ok=True)
+        os.makedirs(deactivate_folder, exist_ok=True)
+
+        activate_script = os.path.join(conda_prefix, activate_folder, f"env_vars.sh")
+        deactivate_script = os.path.join(conda_prefix, deactivate_folder, f"env_vars.sh")
+
+        if os.path.isfile(activate_script):
+            # Only create file if it does not already exist. There may be instances where people
+            # have created their own scripts, but these should be few and far between and those
+            # people should already know what they are doing.
+            return
+
+        conda_libs = os.path.join(conda_prefix, "lib")
+        shebang = "#!/bin/sh\n\n"
+
+        with open(activate_script, "w", encoding="utf8") as afile:
+            afile.write(f"{shebang}")
+            afile.write("export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}\n")
+            afile.write(f"export LD_LIBRARY_PATH='{conda_libs}':${{LD_LIBRARY_PATH}}\n")
+
+        with open(deactivate_script, "w", encoding="utf8") as afile:
+            afile.write(f"{shebang}")
+            afile.write("export LD_LIBRARY_PATH=${OLD_LD_LIBRARY_PATH}\n")
+            afile.write("unset OLD_LD_LIBRARY_PATH\n")
+
+        self.output.info(f"Cuda search path set to '{conda_libs}'")
 
 
 class Output():
@@ -326,14 +372,14 @@ class Output():
         """ Format INFO Text """
         trm = "INFO    "
         if self.term_support_color:
-            trm = "{}INFO   {} ".format(self.green, self.default_color)
+            trm = f"{self.green}INFO   {self.default_color} "
         print(trm + self.__indent_text_block(text))
 
     def warning(self, text):
         """ Format WARNING Text """
         trm = "WARNING "
         if self.term_support_color:
-            trm = "{}WARNING{} ".format(self.yellow, self.default_color)
+            trm = f"{self.yellow}WARNING{self.default_color} "
         print(trm + self.__indent_text_block(text))
 
     def error(self, text):
@@ -341,7 +387,7 @@ class Output():
         global INSTALL_FAILED  # pylint:disable=global-statement
         trm = "ERROR   "
         if self.term_support_color:
-            trm = "{}ERROR  {} ".format(self.red, self.default_color)
+            trm = f"{self.red}ERROR  {self.default_color} "
         print(trm + self.__indent_text_block(text))
         INSTALL_FAILED = True
 
@@ -473,7 +519,7 @@ class CudaCheck():  # pylint:disable=too-few-public-methods
         Initially just calls `nvcc -V` to get the installed version of Cuda currently in use.
         If this fails, drills down to more OS specific checking methods.
         """
-        chk = Popen("nvcc -V", shell=True, stdout=PIPE, stderr=PIPE)
+        with Popen("nvcc -V", shell=True, stdout=PIPE, stderr=PIPE) as chk:
             stdout, stderr = chk.communicate()
         if not stderr:
             version = re.search(r".*release (?P<cuda>\d+\.\d+)",
@@ -524,7 +570,7 @@ class CudaCheck():  # pylint:disable=too-few-public-methods
         if not cudnn_checkfile:
             return
         found = 0
-        with open(cudnn_checkfile, "r") as ofile:
+        with open(cudnn_checkfile, "r", encoding="utf8") as ofile:
             for line in ofile:
                 if line.lower().startswith("#define cudnn_major"):
                     major = line[line.rfind(" ") + 1:].strip()
@@ -553,7 +599,7 @@ class CudaCheck():  # pylint:disable=too-few-public-methods
         chk = os.popen("ldconfig -p | grep -P \"libcudnn.so.\\d+\" | head -n 1").read()
         chk = chk.strip().replace("libcudnn.so.", "")
         if not chk:
-            return list()
+            return []
 
         cudnn_vers = chk[0]
         header_files = [f"cudnn_v{cudnn_vers}.h"] + self._cudnn_header_files
@@ -574,7 +620,7 @@ class CudaCheck():  # pylint:disable=too-few-public-methods
         """
         # TODO A more reliable way of getting the windows location
         if not self.cuda_path:
-            return list()
+            return []
         scandir = os.path.join(self.cuda_path, "include")
         cudnn_checkfiles = [os.path.join(scandir, header) for header in self._cudnn_header_files]
         return cudnn_checkfiles
@@ -703,7 +749,7 @@ class Install():
                 channel = None if len(pkg) != 2 else pkg[1]
                 pkg = pkg[0]
             if version:
-                pkg = "{}{}".format(pkg, ",".join("".join(spec) for spec in version))
+                pkg = f"{pkg}{','.join(''.join(spec) for spec in version)}"
             if self.env.is_conda and not pkg.startswith("git"):
                 if pkg.startswith("tensorflow-gpu"):
                     # From TF 2.4 onwards, Anaconda Tensorflow becomes a mess. The version of 2.5
@@ -762,13 +808,14 @@ class Install():
                 package = f"\"{package}\""
             condaexe.append(package)
 
-        self.output.info("Installing {}".format(package.replace("\"", "")))
+        clean_pkg = package.replace("\"", "")
+        self.output.info(f"Installing {clean_pkg}")
         shell = self.env.os_version[0] == "Windows"
         try:
             if verbose:
                 run(condaexe, check=True, shell=shell)
             else:
-                with open(os.devnull, "w") as devnull:
+                with open(os.devnull, "w", encoding="utf8") as devnull:
                     run(condaexe, stdout=devnull, stderr=devnull, check=True, shell=shell)
         except CalledProcessError:
             if not conda_only:
@@ -811,14 +858,16 @@ class Install():
         pkgs = ["cudatoolkit", "cudnn"]
         shell = self.env.os_version[0] == "Windows"
         for pkg in pkgs:
-            chk = Popen(condaexe + [pkg], shell=shell, stdout=PIPE)
+            with Popen(condaexe + [pkg], shell=shell, stdout=PIPE) as chk:
                 available = [line.split()
-                         for line in chk.communicate()[0].decode(self.env.encoding).splitlines()
+                             for line
+                             in chk.communicate()[0].decode(self.env.encoding).splitlines()
                              if line.startswith(pkg)]
                 compatible = [req for req in available
                               if (pkg == "cudatoolkit" and req[1].startswith(versions[0]))
                               or (pkg == "cudnn" and versions[0] in req[2]
                                   and req[1].startswith(versions[1]))]
+
             candidate = "==".join(sorted(compatible, key=lambda x: x[1])[-1][:2])
             self.conda_installer(candidate, verbose=True, conda_only=True)
 
@@ -830,6 +879,8 @@ class Tips():
 
     def docker_no_cuda(self):
         """ Output Tips for Docker without Cuda """
+
+        path = os.path.dirname(os.path.realpath(__file__))
         self.output.info(
             "1. Install Docker\n"
             "https://www.docker.com/community-edition\n\n"
@@ -839,7 +890,7 @@ class Tips():
             "# without GUI\n"
             "docker run -tid -p 8888:8888 \\ \n"
             "\t--hostname deepfakes-cpu --name deepfakes-cpu \\ \n"
-            "\t-v {path}:/srv \\ \n"
+            f"\t-v {path}:/srv \\ \n"
             "\tdeepfakes-cpu\n\n"
             "# with gui. tools.py gui working.\n"
             "## enable local access to X11 server\n"
@@ -847,7 +898,7 @@ class Tips():
             "## create container\n"
             "nvidia-docker run -tid -p 8888:8888 \\ \n"
             "\t--hostname deepfakes-cpu --name deepfakes-cpu \\ \n"
-            "\t-v {path}:/srv \\ \n"
+            f"\t-v {path}:/srv \\ \n"
             "\t-v /tmp/.X11-unix:/tmp/.X11-unix \\ \n"
             "\t-e DISPLAY=unix$DISPLAY \\ \n"
             "\t-e AUDIO_GID=`getent group audio | cut -d: -f3` \\ \n"
@@ -856,12 +907,13 @@ class Tips():
             "\t-e UID=`id -u` \\ \n"
             "\tdeepfakes-cpu \n\n"
             "4. Open a new terminal to run faceswap.py in /srv\n"
-            "docker exec -it deepfakes-cpu bash".format(
-                path=os.path.dirname(os.path.realpath(__file__))))
+            "docker exec -it deepfakes-cpu bash")
         self.output.info("That's all you need to do with a docker. Have fun.")
 
     def docker_cuda(self):
         """ Output Tips for Docker wit Cuda"""
+
+        path = os.path.dirname(os.path.realpath(__file__))
         self.output.info(
             "1. Install Docker\n"
             "https://www.docker.com/community-edition\n\n"
@@ -875,7 +927,7 @@ class Tips():
             "# without gui \n"
             "docker run -tid -p 8888:8888 \\ \n"
             "\t--hostname deepfakes-gpu --name deepfakes-gpu \\ \n"
-            "\t-v {path}:/srv \\ \n"
+            f"\t-v {path}:/srv \\ \n"
             "\tdeepfakes-gpu\n\n"
             "# with gui.\n"
             "## enable local access to X11 server\n"
@@ -885,7 +937,7 @@ class Tips():
             "## create container\n"
             "nvidia-docker run -tid -p 8888:8888 \\ \n"
             "\t--hostname deepfakes-gpu --name deepfakes-gpu \\ \n"
-            "\t-v {path}:/srv \\ \n"
+            f"\t-v {path}:/srv \\ \n"
             "\t-v /tmp/.X11-unix:/tmp/.X11-unix \\ \n"
             "\t-e DISPLAY=unix$DISPLAY \\ \n"
             "\t-e AUDIO_GID=`getent group audio | cut -d: -f3` \\ \n"
@@ -894,8 +946,7 @@ class Tips():
             "\t-e UID=`id -u` \\ \n"
             "\tdeepfakes-gpu\n\n"
             "6. Open a new terminal to interact with the project\n"
-            "docker exec deepfakes-gpu python /srv/faceswap.py gui\n".format(
-                path=os.path.dirname(os.path.realpath(__file__))))
+            "docker exec deepfakes-gpu python /srv/faceswap.py gui\n")
 
     def macos(self):
         """ Output Tips for macOS"""

tests/__init__.py

@@ -1,7 +0,0 @@
-#!/usr/bin/env python3
-""" Use custom Importer for importing Keras for tests """
-import sys
-from lib.utils import KerasFinder
-
-
-sys.meta_path.insert(0, KerasFinder())

tests/lib/model/initializers_test.py

@@ -4,14 +4,21 @@
 Adapted from Keras tests.
 """
 
-from keras import backend as K
-from keras import initializers as k_initializers
 import pytest
 import numpy as np
 
 from lib.model import initializers
 from lib.utils import get_backend
 
+if get_backend() == "amd":
+    from keras import backend as K
+    from keras import initializers as k_initializers
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import backend as K  # pylint:disable=import-error
+    from tensorflow.keras import initializers as k_initializers  # pylint:disable=import-error
+
+
 CONV_SHAPE = (3, 3, 256, 2048)
 CONV_ID = get_backend().upper()
 

tests/lib/model/layers_test.py

@@ -7,7 +7,6 @@ Adapted from Keras tests.
 
 import pytest
 import numpy as np
-from keras import Input, Model, backend as K
 
 from numpy.testing import assert_allclose
 
@@ -15,6 +14,13 @@ from lib.model import layers
 from lib.utils import get_backend
 from tests.utils import has_arg
 
+if get_backend() == "amd":
+    from keras import Input, Model, backend as K
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import Input, Model, backend as K  # pylint:disable=import-error
+
+
 CONV_SHAPE = (3, 3, 256, 2048)
 CONV_ID = get_backend().upper()
 

tests/lib/model/losses_test.py

@@ -6,17 +6,15 @@ Adapted from Keras tests.
 
 import pytest
 import numpy as np
-from numpy.testing import assert_allclose
-
-from keras import backend as K
-from keras import losses as k_losses
-from keras.layers import Conv2D
-from keras.models import Sequential
-from keras.optimizers import Adam
 
 from lib.model import losses
 from lib.utils import get_backend
 
+if get_backend() == "amd":
+    from keras import backend as K, losses as k_losses
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import backend as K, losses as k_losses  # pylint:disable=import-error
 
 _PARAMS = [(losses.GeneralizedLoss(), (2, 16, 16)),
            (losses.GradientLoss(), (2, 16, 16)),
@@ -25,7 +23,7 @@ _PARAMS = [(losses.GeneralizedLoss(), (2, 16, 16)),
            # TODO Make sure these output dimensions are correct
            (losses.LInfNorm(), (2, 1, 1))]
 _IDS = ["GeneralizedLoss", "GradientLoss", "GMSDLoss", "LInfNorm"]
-_IDS = ["{}[{}]".format(loss, get_backend().upper()) for loss in _IDS]
+_IDS = [f"{loss}[{get_backend().upper()}]" for loss in _IDS]
 
 
 @pytest.mark.parametrize(["loss_func", "output_shape"], _PARAMS, ids=_IDS)
@@ -45,10 +43,10 @@ def test_loss_output(loss_func, output_shape):
 
 _LWPARAMS = [losses.GeneralizedLoss(), losses.GradientLoss(), losses.GMSDLoss(),
              losses.LInfNorm(), k_losses.mean_absolute_error, k_losses.mean_squared_error,
-             k_losses.logcosh, losses.DSSIMObjective()]
+             k_losses.logcosh, losses.DSSIMObjective(), losses.MSSSIMLoss()]
 _LWIDS = ["GeneralizedLoss", "GradientLoss", "GMSDLoss", "LInfNorm", "mae", "mse", "logcosh",
-          "DSSIMObjective"]
-_LWIDS = ["{}[{}]".format(loss, get_backend().upper()) for loss in _LWIDS]
+          "DSSIMObjective", "MS-SSIM"]
+_LWIDS = [f"{loss}[{get_backend().upper()}]" for loss in _LWIDS]
 
 
 @pytest.mark.parametrize("loss_func", _LWPARAMS, ids=_LWIDS)
@@ -57,6 +55,8 @@ def test_loss_wrapper(loss_func):
     if get_backend() == "amd":
         if isinstance(loss_func, losses.GMSDLoss):
             pytest.skip("GMSD Loss is not currently compatible with PlaidML")
+        if isinstance(loss_func, losses.MSSSIMLoss):
+            pytest.skip("MS-SSIM Loss is not currently compatible with PlaidML")
         if hasattr(loss_func, "__name__") and loss_func.__name__ == "logcosh":
             pytest.skip("LogCosh Loss is not currently compatible with PlaidML")
     y_a = K.variable(np.random.random((2, 16, 16, 4)))
@@ -70,41 +70,3 @@ def test_loss_wrapper(loss_func):
     else:
         output = output.numpy()
         assert output.dtype == "float32" and not np.isnan(output)
-
-
-@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
-def test_dssim_channels_last(dummy):  # pylint:disable=unused-argument
-    """ Basic test for DSSIM Loss """
-    prev_data = K.image_data_format()
-    K.set_image_data_format('channels_last')
-    for input_dim, kernel_size in zip([32, 33], [2, 3]):
-        input_shape = [input_dim, input_dim, 3]
-        var_x = np.random.random_sample(4 * input_dim * input_dim * 3)
-        var_x = var_x.reshape([4] + input_shape)
-        var_y = np.random.random_sample(4 * input_dim * input_dim * 3)
-        var_y = var_y.reshape([4] + input_shape)
-
-        model = Sequential()
-        model.add(Conv2D(32, (3, 3), padding='same', input_shape=input_shape,
-                         activation='relu'))
-        model.add(Conv2D(3, (3, 3), padding='same', input_shape=input_shape,
-                         activation='relu'))
-        adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-8)
-        model.compile(loss=losses.DSSIMObjective(kernel_size=kernel_size),
-                      metrics=['mse'],
-                      optimizer=adam)
-        model.fit(var_x, var_y, batch_size=2, epochs=1, shuffle='batch')
-
-        # Test same
-        x_1 = K.constant(var_x, 'float32')
-        x_2 = K.constant(var_x, 'float32')
-        dssim = losses.DSSIMObjective(kernel_size=kernel_size)
-        assert_allclose(0.0, K.eval(dssim(x_1, x_2)), atol=1e-4)
-
-        # Test opposite
-        x_1 = K.zeros([4] + input_shape)
-        x_2 = K.ones([4] + input_shape)
-        dssim = losses.DSSIMObjective(kernel_size=kernel_size)
-        assert_allclose(0.5, K.eval(dssim(x_1, x_2)), atol=1e-4)
-
-    K.set_image_data_format(prev_data)

tests/lib/model/nn_blocks_test.py

@@ -9,12 +9,18 @@ from itertools import product
 import pytest
 import numpy as np
 
-from keras import Input, Model, backend as K
+
 from numpy.testing import assert_allclose
 
 from lib.model import nn_blocks
 from lib.utils import get_backend
 
+if get_backend() == "amd":
+    from keras import Input, Model, backend as K
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import Input, Model, backend as K  # pylint:disable=import-error
+
 
 def block_test(layer_func, kwargs={}, input_shape=None):
     """Test routine for faceswap neural network blocks.

tests/lib/model/normalization_test.py

@@ -8,13 +8,17 @@ from itertools import product
 import numpy as np
 import pytest
 
-from keras import regularizers, models, layers
-
 from lib.model import normalization
 from lib.utils import get_backend
 
 from tests.lib.model.layers_test import layer_test
 
+if get_backend() == "amd":
+    from keras import regularizers, models, layers
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import regularizers, models, layers  # pylint:disable=import-error
+
 
 @pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
 def test_instance_normalization(dummy):  # pylint:disable=unused-argument
@@ -101,8 +105,7 @@ def test_layer_normalization(center, scale):
 
 _PARAMS = ["partial", "bias"]
 _VALUES = [(0.0, False), (0.25, False), (0.5, True), (0.75, False), (1.0, True)]
-_IDS = ["partial={}|bias={}[{}]".format(v[0], v[1], get_backend().upper())
-        for v in _VALUES]
+_IDS = [f"partial={v[0]}|bias={v[1]}[{get_backend().upper()}]" for v in _VALUES]
 
 
 @pytest.mark.parametrize(_PARAMS, _VALUES, ids=_IDS)

tests/lib/model/optimizers_test.py

@@ -5,9 +5,6 @@ Adapted from Keras tests.
 """
 import pytest
 
-from keras import optimizers as k_optimizers
-from keras.layers import Dense, Activation
-from keras.models import Sequential
 import numpy as np
 from numpy.testing import assert_allclose
 
@@ -16,6 +13,16 @@ from lib.utils import get_backend
 
 from tests.utils import generate_test_data, to_categorical
 
+if get_backend() == "amd":
+    from keras import optimizers as k_optimizers
+    from keras.layers import Dense, Activation
+    from keras.models import Sequential
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow.keras import optimizers as k_optimizers  # pylint:disable=import-error
+    from tensorflow.keras.layers import Dense, Activation  # noqa pylint:disable=import-error,no-name-in-module
+    from tensorflow.keras.models import Sequential  # pylint:disable=import-error,no-name-in-module
+
 
 def get_test_data():
     """ Obtain randomized test data for training """
@@ -76,11 +83,11 @@ def _test_optimizer(optimizer, target=0.75):
 @pytest.mark.parametrize("dummy", [None], ids=[get_backend().upper()])
 def test_adam(dummy):  # pylint:disable=unused-argument
     """ Test for custom Adam optimizer """
-    _test_optimizer(k_optimizers.Adam(), target=0.5)
-    _test_optimizer(k_optimizers.Adam(decay=1e-3), target=0.5)
+    _test_optimizer(k_optimizers.Adam(), target=0.45)
+    _test_optimizer(k_optimizers.Adam(decay=1e-3), target=0.45)
 
 
 @pytest.mark.parametrize("dummy", [None], ids=[get_backend().upper()])
 def test_adabelief(dummy):  # pylint:disable=unused-argument
     """ Test for custom Adam optimizer """
-    _test_optimizer(optimizers.AdaBelief(), target=0.5)
+    _test_optimizer(optimizers.AdaBelief(), target=0.45)

tests/startup_test.py

@@ -5,10 +5,16 @@ import inspect
 
 import pytest
 
+from lib.utils import get_backend
+
+if get_backend() == "amd":
     import keras
     from keras import backend as K
+else:
+    # Ignore linting errors from Tensorflow's thoroughly broken import system
+    from tensorflow import keras
+    from tensorflow.keras import backend as K  # pylint:disable=import-error
 
-from lib.utils import get_backend
 
 _BACKEND = get_backend()
 
@@ -25,5 +31,6 @@ def test_backend(dummy):  # pylint:disable=unused-argument
 def test_keras(dummy):  # pylint:disable=unused-argument
     """ Sanity check to ensure that tensorflow keras is being used for CPU and standard
     keras for AMD. """
-    assert ((_BACKEND == "cpu" and keras.__version__ in ("2.3.0-tf", "2.4.0")) or
+    assert ((_BACKEND == "cpu" and keras.__version__ in ("2.3.0-tf", "2.4.0",
+                                                         "2.6.0", "2.7.0", "2.8.0")) or
             (_BACKEND == "amd" and keras.__version__ == "2.2.4"))

tools/alignments/jobs.py

@@ -405,13 +405,30 @@ class Extract():  # pylint:disable=too-few-public-methods
         self._is_legacy = self._alignments.version == 1.0  # pylint:disable=protected-access
         self._mask_pipeline = None
         self._faces_dir = arguments.faces_dir
-        self._frames = Frames(arguments.frames_dir)
+
+        self._frames = Frames(arguments.frames_dir, self._get_count())
         self._extracted_faces = ExtractedFaces(self._frames,
                                                self._alignments,
                                                size=arguments.size)
         self._saver = None
         logger.debug("Initialized %s", self.__class__.__name__)
 
+    def _get_count(self):
+        """ If the alignments file has been run through the manual tool, then it will hold video
+        meta information, meaning that the count of frames in the alignment file can be relied
+        on to be accurate.
+
+        Returns
+        -------
+        int or ``None``
+        For video input which contain video meta-data in the alignments file then the count of
+        frames is returned. In all other cases ``None`` is returned
+        """
+        has_meta = all(val is not None for val in self._alignments.video_meta_data.values())
+        retval = len(self._alignments.video_meta_data["pts_time"]) if has_meta else None
+        logger.debug("Frame count from alignments file: (has_meta: %s, %s", has_meta, retval)
+        return retval
+
     def process(self):
         """ Run the re-extraction from Alignments file process"""
         logger.info("[EXTRACT FACES]")  # Tidy up cli output

tools/alignments/media.py

@@ -72,11 +72,14 @@ class MediaLoader():
     ----------
     folder: str
         The folder of images or video file to load images from
+    count: int or ``None``, optional
+        If the total frame count is known it can be passed in here which will skip
+        analyzing a video file. If the count is not passed in, it will be calculated.
     """
-    def __init__(self, folder):
+    def __init__(self, folder, count=None):
         logger.debug("Initializing %s: (folder: '%s')", self.__class__.__name__, folder)
         logger.info("[%s DATA]", self.__class__.__name__.upper())
-        self._count = None
+        self._count = count
         self.folder = folder
         self.vid_reader = self.check_input_folder()
         self.file_list_sorted = self.sorted_items()
@@ -188,7 +191,7 @@ class MediaLoader():
         numpy.ndarray
             The image that has been loaded from disk
         """
-        loader = ImagesLoader(self.folder, queue_size=32)
+        loader = ImagesLoader(self.folder, queue_size=32, count=self._count)
         if skip_list is not None:
             loader.add_skip_list(skip_list)
         for filename, image in loader.load():

tools/manual/detected_faces.py

@@ -255,7 +255,9 @@ class _DiskIO():  # pylint:disable=too-few-public-methods
         self._tk_edited = detected_faces.tk_edited
         self._tk_face_count_changed = detected_faces.tk_face_count_changed
         self._globals = detected_faces._globals
-        self._sorted_frame_names = sorted(self._alignments.data)
+
+        # Must be populated after loading faces as video_meta_data may have increased frame count
+        self._sorted_frame_names = None
         logger.debug("Initialized %s", self.__class__.__name__)
 
     def load(self):
@@ -270,6 +272,7 @@ class _DiskIO():  # pylint:disable=too-few-public-methods
                 _ = face.aligned.average_distance  # cache the distances
                 this_frame_faces.append(face)
             self._frame_faces.append(this_frame_faces)
+        self._sorted_frame_names = sorted(self._alignments.data)
 
     def save(self):
         """ Convert updated :class:`~lib.align.DetectedFace` objects to alignments format

tools/sort/cli.py

@@ -146,7 +146,7 @@ class SortArgs(FaceSwapArgs):
                    "right. If the number of images doesn't divide evenly into the number of "
                    "bins, the remaining images get put in the last bin. For black-pixels it "
                    "represents the divider of the percentage of black pixels. For 10, first "
-                   "folder will have the faces with 0 to 10% black pixels, second 11 to 20%, "
+                   "folder will have the faces with 0 to 10%% black pixels, second 11 to 20%%, "
                    "etc. Default value: 5")))
         argument_list.append(dict(
             opts=('-l', '--log-changes'),