Update extract pipeline

plugins/extract/_base.py

@@ -9,11 +9,10 @@
 from dataclasses import dataclass, field
 
 import numpy as np
-from tensorflow.python.framework import errors_impl as tf_errors  # pylint:disable=no-name-in-module  # noqa
 
 from lib.multithreading import MultiThread
 from lib.queue_manager import queue_manager
-from lib.utils import GetModel, FaceswapError
+from lib.utils import GetModel
 from ._config import Config
 from .pipeline import ExtractMedia
 
@@ -109,7 +108,7 @@ class Extractor():
     model_filename: str
         The name of the model file to be loaded
     exclude_gpus: list, optional
-        A list of indices correlating to connected GPUs that Tensorflow should not use. Pass
+        A list of indices correlating to connected GPUs that PyTorch should not use. Pass
         ``None`` to not exclude any GPUs. Default: ``None``
     configfile: str, optional
         Path to a custom configuration ``ini`` file. Default: Use system configfile
@@ -135,9 +134,6 @@ class Extractor():
     vram: int
         Approximate VRAM used by the model at :attr:`input_size`. Used to calculate the
         :attr:`batchsize`. Be conservative to avoid OOM.
-    vram_warnings: int
-        Approximate VRAM used by the model at :attr:`input_size` that will still run, but generates
-        warnings. Used to calculate the :attr:`batchsize`. Be conservative to avoid OOM.
     vram_per_batch: int
         Approximate additional VRAM used by the model for each additional batch. Used to calculate
         the :attr:`batchsize`. Be conservative to avoid OOM.
@@ -174,7 +170,6 @@ def __init__(self,
         self.input_size = 0
         self.color_format: T.Literal["BGR", "RGB", "GRAY"] = "BGR"
         self.vram = 0
-        self.vram_warnings = 0  # Will run at this with warnings
         self.vram_per_batch = 0
 
         # << THE FOLLOWING ARE SET IN self.initialize METHOD >> #
@@ -470,20 +465,7 @@ def initialize(self, *args, **kwargs) -> None:
                          kwargs["out_queue"],
                          [f"predict_{name}", f"post_{name}"])
         self._compile_threads()
-        try:
-            self.init_model()
-        except tf_errors.UnknownError as err:
-            if "failed to get convolution algorithm" in str(err).lower():
-                msg = ("Tensorflow raised an unknown error. This is most likely caused by a "
-                       "failure to launch cuDNN which can occur for some GPU/Tensorflow "
-                       "combinations. You should enable `allow_growth` to attempt to resolve this "
-                       "issue:"
-                       "\nGUI: Go to Settings > Extract Plugins > Global and enable the "
-                       "`allow_growth` option."
-                       "\nCLI: Go to `faceswap/config/extract.ini` and change the `allow_growth "
-                       "option to `True`.")
-                raise FaceswapError(msg) from err
-            raise err
+        self.init_model()
         self._is_initialized = True
         logger.info("Initialized %s (%s) with batchsize of %s",
                     self.name, self._plugin_type.title(), self.batchsize)
@@ -598,20 +580,7 @@ def _thread_process(self,
                 break
             if not batch.filename:  # Batch not populated. Possible during re-aligns
                 continue
-            try:
-                batch = function(batch)
-            except tf_errors.UnknownError as err:
-                if "failed to get convolution algorithm" in str(err).lower():
-                    msg = ("Tensorflow raised an unknown error. This is most likely caused by a "
-                           "failure to launch cuDNN which can occur for some GPU/Tensorflow "
-                           "combinations. You should enable `allow_growth` to attempt to resolve "
-                           "this issue:"
-                           "\nGUI: Go to Settings > Extract Plugins > Global and enable the "
-                           "`allow_growth` option."
-                           "\nCLI: Go to `faceswap/config/extract.ini` and change the "
-                           "`allow_growth option to `True`.")
-                    raise FaceswapError(msg) from err
-                raise err
+            batch = function(batch)
             if function.__name__ == "_process_output":
                 # Process output items to individual items from batch
                 for item in self.finalize(batch):

plugins/extract/_config.py

@@ -30,16 +30,6 @@ def set_globals(self) -> None:
         logger.debug("Setting global config")
         section = "global"
         self.add_section(section, _("Options that apply to all extraction plugins"))
-        self.add_item(
-            section=section,
-            title="allow_growth",
-            datatype=bool,
-            default=False,
-            group=_("settings"),
-            info=_("Enable the Tensorflow GPU `allow_growth` configuration option. "
-                   "This option prevents Tensorflow from allocating all of the GPU VRAM at launch "
-                   "but can lead to higher VRAM fragmentation and slower performance. Should only "
-                   "be enabled if you are having problems running extraction."))
         self.add_item(
             section=section,
             title="aligner_min_scale",

plugins/extract/align/_base/aligner.py

@@ -21,8 +21,7 @@
 
 import cv2
 import numpy as np
-
-from tensorflow.python.framework import errors_impl as tf_errors  # pylint:disable=no-name-in-module # noqa
+from torch.cuda import OutOfMemoryError
 
 from lib.utils import FaceswapError
 from plugins.extract._base import BatchType, Extractor, ExtractMedia, ExtractorBatch
@@ -534,6 +533,7 @@ def _predict(self, batch: BatchType) -> AlignerBatch:
             preds = [self.predict(feed) for feed in batch.refeeds]
             try:
                 batch.prediction = np.array(preds)
+                logger.trace("Aligner out: %s", batch.prediction.shape)
             except ValueError as err:
                 # If refeed batches are different sizes, Numpy will error, so we need to explicitly
                 # set the dtype to 'object' rather than let it infer
@@ -549,8 +549,7 @@ def _predict(self, batch: BatchType) -> AlignerBatch:
                 else:
                     raise
 
-            return batch
-        except tf_errors.ResourceExhaustedError as err:
+        except OutOfMemoryError as err:
             msg = ("You do not have enough GPU memory available to run detection at the "
                    "selected batch size. You can try a number of things:"
                    "\n1) Close any other application that is using your GPU (web browsers are "
@@ -561,6 +560,8 @@ def _predict(self, batch: BatchType) -> AlignerBatch:
                    "\n3) Enable 'Single Process' mode.")
             raise FaceswapError(msg) from err
 
+        return batch
+
     def _process_refeeds(self, batch: AlignerBatch) -> list[AlignerBatch]:
         """ Process the output for each selected re-feed
 

plugins/extract/align/fan.py

@@ -29,9 +29,8 @@ def __init__(self, **kwargs) -> None:
         self.name = "FAN"
         self.input_size = 256
         self.color_format = "RGB"
-        self.vram = 2240
-        self.vram_warnings = 512  # Will run at this with warnings
-        self.vram_per_batch = 64
+        self.vram = 896  # 810 in testing
+        self.vram_per_batch = 768  # ~720 in testing
         self.realign_centering = "head"
         self.batchsize: int = self.config["batch-size"]
         self.reference_scale = 200. / 195.
@@ -42,7 +41,6 @@ def init_model(self) -> None:
         assert isinstance(self.model_path, str)
         self.model = KSession(self.name,
                               self.model_path,
-                              allow_growth=self.config["allow_growth"],
                               exclude_gpus=self._exclude_gpus)
         self.model.load_model()
         # Feed a placeholder so Aligner is primed for Manual tool
@@ -224,7 +222,6 @@ def predict(self, feed: np.ndarray) -> np.ndarray:
         # TODO Remove lazy transpose and change points from predict to use the correct
         # order
         retval = self.model.predict(feed)[-1].transpose(0, 3, 1, 2)
-        logger.trace(retval.shape)  # type:ignore[attr-defined]
         return retval
 
     def process_output(self, batch: BatchType) -> None:

plugins/extract/detect/_base.py

@@ -23,8 +23,7 @@
 
 import cv2
 import numpy as np
-
-from tensorflow.python.framework import errors_impl as tf_errors  # pylint:disable=no-name-in-module # noqa
+from torch.cuda import OutOfMemoryError
 
 from lib.align import DetectedFace
 from lib.utils import FaceswapError
@@ -280,15 +279,7 @@ def _predict(self, batch: BatchType) -> DetectorBatch:
             self._rotate_batch(batch, angle)
             try:
                 pred = self.predict(batch.feed)
-                if angle == 0:
-                    batch.prediction = pred
-                else:
-                    batch.prediction = np.array([b if b.any() else p
-                                                 for b, p in zip(batch.prediction, pred)])
-                logger.trace("angle: %s, filenames: %s, "  # type:ignore[attr-defined]
-                             "prediction: %s",
-                             angle, batch.filename, pred)
-            except tf_errors.ResourceExhaustedError as err:
+            except OutOfMemoryError as err:
                 msg = ("You do not have enough GPU memory available to run detection at the "
                        "selected batch size. You can try a number of things:"
                        "\n1) Close any other application that is using your GPU (web browsers are "
@@ -299,6 +290,15 @@ def _predict(self, batch: BatchType) -> DetectorBatch:
                        "\n3) Enable 'Single Process' mode.")
                 raise FaceswapError(msg) from err
 
+            if angle == 0:
+                batch.prediction = pred
+            else:
+                batch.prediction = np.array([b if b.any() else p
+                                                for b, p in zip(batch.prediction, pred)])
+            logger.trace("angle: %s, filenames: %s, "  # type:ignore[attr-defined]
+                            "prediction: %s",
+                            angle, batch.filename, pred)
+
             if angle != 0 and any(face.any() for face in batch.prediction):
                 logger.verbose("found face(s) by rotating image %s "  # type:ignore[attr-defined]
                                "degrees",