tools/model/model.py - master - githubmirror/faceswap

tools/model/model.py
280 строк9 KB
torzdf
bugfix: Inference model. Make more robust to esoteric layouts
27 мар 2026, 12:36
27 мар 2026, 12:36c8abd1b
#!/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__)