/
githubmirror
/
faceswap
Обзор
Документация
Войти
/
githubmirror
/
faceswap
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
lib/gpu_stats/apple_silicon.py
191 строка
6 KB
torzdf
Backported code to support future updates
27 фев 2026, 22:30
27 фев 2026, 22:30
aa40efa
Код
Авторство
О чём код?
#!/usr/bin/env python3 """Collects and returns Information on available Apple Silicon SoCs in Apple Macs.""" import typing as T import os import psutil import torch from lib.utils import FaceswapError, get_module_objects from ._base import _GPUStats _METAL_INITIALIZED: bool = False class AppleSiliconStats(_GPUStats): """Holds information and statistics about Apple Silicon SoC(s) available on the currently running Apple system. Notes ----- Apple Silicon is a bit different from other backends, as it does not have a dedicated GPU with it's own dedicated VRAM, rather the RAM is shared with the CPU and GPU. A combination of psutil and torch are used to pull as much useful information as possible. Parameters ---------- log Whether the class should output information to the logger. There may be occasions where the logger has not yet been set up when this class is queried. Attempting to log in these instances will raise an error. If GPU stats are being queried prior to the logger being available then this parameter should be set to ``False``. Otherwise set to ``True``. Default: ``True`` """ def __init__(self, log: bool = True) -> None: # Following attribute set in :func:``_initialize`` self._mps_devices: list[T.Any] = [] super().__init__(log=log) def _initialize(self) -> None: """Initialize Metal for Apple Silicon SoC(s). If :attr:`_is_initialized` is ``True`` then this function just returns performing no action. Otherwise :attr:`is_initialized` is set to ``True`` after successfully initializing Metal. """ if self._is_initialized: return self._log("debug", "Initializing Metal for Apple Silicon SoC.") self._initialize_metal() self._mps_devices = [torch.device("mps")] super()._initialize() def _initialize_metal(self) -> None: """Initialize Metal on first call to this class and set global :attr:``_METAL_INITIALIZED`` to ``True``. If Metal has already been initialized then return performing no action.""" global _METAL_INITIALIZED # pylint:disable=global-statement if _METAL_INITIALIZED: return self._log("debug", "Performing first time Apple SoC setup.") os.environ["DISPLAY"] = ":0" try: os.system("open -a XQuartz") except Exception as err: # pylint:disable=broad-except self._log("debug", f"Swallowing error opening XQuartz: {str(err)}") self._test_torch() _METAL_INITIALIZED = True def _test_torch(self) -> None: """Test that torch can execute correctly. Raises ------ FaceswapError If the Torch library could not be successfully initialized """ try: meminfo = torch.mps.driver_allocated_memory() self._log("debug", f"Torch initialization test: (mem_info: {meminfo})") except RuntimeError as err: msg = ("An unhandled exception occurred initializing the device via Torch " f"Library. Original error: {str(err)}") raise FaceswapError(msg) from err def _get_device_count(self) -> int: """Detect the number of SoCs attached to the system. Returns ------- The total number of SoCs available """ retval = len(self._mps_devices) self._log("debug", f"GPU Device count: {retval}") return retval def _get_handles(self) -> list: """Obtain the device handles for all available Apple Silicon SoCs. Notes ----- Apple SoC does not use handles, so return a list of indices corresponding to found GPU devices Returns ------- The list of indices for available Apple Silicon SoCs """ handles = list(range(self._device_count)) self._log("debug", f"GPU Handles found: {handles}") return handles def _get_driver(self) -> str: """Obtain the Apple Silicon driver version currently in use. Notes ----- As the SoC is not a discreet GPU it does not technically have a driver version, so just return `'Not Applicable'` as a string Returns ------- The current SoC driver version """ driver = "Not Applicable" self._log("debug", f"GPU Driver: {driver}") return driver def _get_device_names(self) -> list[str]: """Obtain the list of names of available Apple Silicon SoC(s) as identified in :attr:`_handles`. Returns ------- The list of available Apple Silicon SoC names """ names = [d.type for d in self._mps_devices] self._log("debug", f"GPU Devices: {names}") return names def _get_vram(self) -> list[int]: """Obtain the VRAM in Megabytes for each available Apple Silicon SoC(s) as identified in :attr:`_handles`. Returns ------- The RAM in Megabytes for each available Apple Silicon SoC """ vram = [int((torch.mps.driver_allocated_memory() / self._device_count) / (1024 * 1024)) for _ in range(self._device_count)] self._log("debug", f"SoC RAM: {vram}") return vram def _get_free_vram(self) -> list[int]: """Obtain the amount of VRAM that is available, in Megabytes, for each available Apple Silicon SoC. Returns ------- List of `float`s containing the amount of RAM available, in Megabytes, for each available SoC as corresponding to the values in :attr:`_handles """ vram = [int((psutil.virtual_memory().available / self._device_count) / (1024 * 1024)) for _ in range(self._device_count)] self._log("debug", f"SoC RAM free: {vram}") return vram def exclude_devices(self, devices: list[int]) -> None: """Apple-Silicon does not support excluding devices Parameters ---------- devices The GPU device IDS to be excluded """ self._log("warning", "Apple Silicon does not support excluding GPUs. This option has been " "ignored") __all__ = get_module_objects(__name__)