GithubStar/faceswap
1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap.git synced 2026-10-17 05:21:10 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
master
faceswap/tools/model/model.py
T

281 lines
9.5 KiB
Python
Raw Permalink Blame History

#!/usr/bin/env python3
"""Tool to restore models from backup"""
from __future__ import annotations
import logging
import os
import sys
import typing as T
from keras import saving
import numpy as np
import keras
from lib.model.backup_restore import Backup
from lib.logger import parse_class_init
# Import the following libs for custom objects
from lib.model import initializers, layers, normalization # noqa # pylint:disable=unused-import
from lib.utils import get_module_objects
from plugins.train.model._base.model import Inference as FSInference
if T.TYPE_CHECKING:
import argparse
logger = logging.getLogger(__name__)
class Model():
"""Tool to perform actions on a model file.
Parameters
----------
arguments
The command line arguments calling the model tool
"""
def __init__(self, arguments: argparse.Namespace) -> None:
logger.debug(parse_class_init(locals()))
self._model_dir = self._check_folder(arguments.model_dir)
self._job = self._get_job(arguments)
logger.debug("Initialized %s", self.__class__.__name__)
@classmethod
def _get_job(cls, arguments: argparse.Namespace) -> Inference | NaNScan | Restore:
"""Get the correct object that holds the selected job.
Parameters
----------
arguments
The command line arguments received for the Model tool which will be used to initiate
the selected job
Returns
-------
The object that will perform the selected job
"""
jobs: dict[str, T.Type[Inference | NaNScan | Restore]] = {
"inference": Inference,
"nan-scan": NaNScan,
"restore": Restore}
return jobs[arguments.job](arguments)
@classmethod
def _check_folder(cls, model_dir: str) -> str:
"""Check that the passed in model folder exists and contains a valid model.
If the passed in value fails any checks, process exits.
Parameters
----------
model_dir
The model folder to be checked
Returns
-------
The confirmed location of the model folder.
"""
if not os.path.exists(model_dir):
logger.error("Model folder does not exist: '%s'", model_dir)
sys.exit(1)
chk_files = [fname
for fname in os.listdir(model_dir)
if fname.endswith(".keras")
and not os.path.splitext(fname)[0].endswith("_inference")]
if not chk_files:
logger.error("Could not find a model in the supplied folder: '%s'", model_dir)
sys.exit(1)
if len(chk_files) > 1:
logger.error("More than one model file found in the model folder: '%s'", model_dir)
sys.exit(1)
model_name = os.path.splitext(chk_files[0])[0].title()
logger.info("%s Model found", model_name)
return model_dir
def process(self) -> None:
"""Call the selected model job."""
self._job.process()
class Inference():
"""Save an inference model from a trained Faceswap model.
Parameters
----------
arguments
The command line arguments calling the model tool
"""
def __init__(self, arguments: argparse.Namespace) -> None:
logger.debug(parse_class_init(locals()))
self._switch = arguments.swap_model
self._input_file, self._output_file = self._get_output_file(arguments.model_dir)
logger.debug("Initialized %s", self.__class__.__name__)
def _get_output_file(self, model_dir: str) -> tuple[str, str]:
"""Obtain the full path for the output model file/folder
Parameters
----------
model_dir
The full path to the folder containing the Faceswap trained model .keras file
Returns
-------
source_model
The full path to the source model file
inference_model
The full path to the inference model save location
"""
model_name = next(fname for fname in os.listdir(model_dir)
if fname.endswith(".keras")
and not fname.endswith("_inference.keras"))
in_path = os.path.join(model_dir, model_name)
logger.debug("Model input path: '%s'", in_path)
model_name = f"{os.path.splitext(model_name)[0]}_inference.keras"
out_path = os.path.join(model_dir, model_name)
logger.debug("Inference output path: '%s'", out_path)
return in_path, out_path
def process(self) -> None:
"""Run the inference model creation process."""
logger.info("Loading model '%s'", self._input_file)
model = saving.load_model(self._input_file, compile=False)
logger.info("Creating inference model...")
inference = FSInference(model, self._switch)()
logger.info("Saving to: '%s'", self._output_file)
inference.save(self._output_file)
class NaNScan():
"""Tool to scan for NaN and Infs in model weights.
Parameters
----------
arguments
The command line arguments calling the model tool
"""
def __init__(self, arguments: argparse.Namespace) -> None:
logger.debug(parse_class_init(locals()))
self._model_file = self._get_model_filename(arguments.model_dir)
logger.debug("Initialized %s", self.__class__.__name__)
@classmethod
def _get_model_filename(cls, model_dir: str) -> str:
"""Obtain the full path the model's .keras file.
Parameters
----------
model_dir
The full path to the folder containing the model file
Returns
-------
The full path to the saved model file
"""
model_file = next(fname for fname in os.listdir(model_dir) if fname.endswith(".keras"))
return os.path.join(model_dir, model_file)
def _parse_weights(self,
layer: keras.models.Model | keras.layers.Layer) -> dict:
"""Recursively pass through sub-models to scan layer weights"""
weights = layer.get_weights()
logger.debug("Processing weights for layer '%s', length: '%s'",
layer.name, len(weights))
if not weights:
logger.debug("Skipping layer with no weights: %s", layer.name)
return {}
if hasattr(layer, "layers"): # Must be a sub-model
retval = {}
for lyr in layer.layers:
info = self._parse_weights(lyr)
if not info:
continue
retval[lyr.name] = info
return retval
nans = sum(np.count_nonzero(np.isnan(w)) for w in weights)
infs = sum(np.count_nonzero(np.isinf(w)) for w in weights)
if nans + infs == 0:
return {}
return {"nans": nans, "infs": infs}
def _parse_output(self, errors: dict, indent: int = 0) -> None:
"""Parse the output of the errors dictionary and print a pretty summary.
Parameters
----------
errors
The nested dictionary of errors found when parsing the weights
indent
How far should the current printed line be indented. Default: `0`
"""
for key, val in errors.items():
log_line = f"|{'--' * indent} "
log_line += key.ljust(50 - len(log_line))
if isinstance(val, dict) and "nans" not in val:
logger.info(log_line)
self._parse_output(val, indent + 1)
elif isinstance(val, dict) and "nans" in val:
log_line += f"nans: {val['nans']}, infs: {val['infs']}"
logger.info(log_line.ljust(30))
def process(self) -> None:
"""Scan the loaded model for NaNs and Infs and output summary."""
logger.info("Loading model...")
model = saving.load_model(self._model_file, compile=False)
logger.info("Parsing weights for invalid values...")
errors = self._parse_weights(model)
if not errors:
logger.info("No invalid values found in model: '%s'", self._model_file)
sys.exit(1)
logger.info("Invalid values found in model: %s", self._model_file)
self._parse_output(errors)
class Restore():
"""Restore a model from backup.
Parameters
----------
arguments
The command line arguments calling the model tool
"""
def __init__(self, arguments: argparse.Namespace) -> None:
logger.debug(parse_class_init(locals()))
self._model_dir = arguments.model_dir
self._model_name = self._get_model_name()
logger.debug("Initialized %s", self.__class__.__name__)
def process(self) -> None:
"""Perform the Restore process"""
logger.info("Starting Model Restore...")
backup = Backup(self._model_dir, self._model_name)
backup.restore()
logger.info("Completed Model Restore")
def _get_model_name(self) -> str:
"""Additional checks to make sure that a backup exists in the model location."""
bk_files = [fname for fname in os.listdir(self._model_dir) if fname.endswith(".bk")]
if not bk_files:
logger.error("Could not find any backup files in the supplied folder: '%s'",
self._model_dir)
sys.exit(1)
logger.verbose("Backup files: %s)", bk_files) # type:ignore[attr-defined]
ext = ".keras.bk"
model_name = next(fname for fname in bk_files if fname.endswith(ext))
return model_name[:-len(ext)]
__all__ = get_module_objects(__name__)
Reference in New Issue View Git Blame Copy Permalink