diff --git a/facefusion/execution.py b/facefusion/execution.py index 213d0ced..2d768c16 100644 --- a/facefusion/execution.py +++ b/facefusion/execution.py @@ -8,6 +8,7 @@ import pynvml import onnxruntime import facefusion.choices +from facefusion import state_manager from facefusion.system import detect_static_graphic_devices from facefusion.types import ExecutionProvider, InferenceSessionProvider @@ -113,7 +114,7 @@ def resolve_cache_path() -> str: def resolve_cudnn_conv_algo_search() -> str: - execution_devices = detect_static_graphic_devices() + execution_devices = detect_static_graphic_devices(state_manager.get_item('execution_providers')) product_names = ('GeForce GTX 1630', 'GeForce GTX 1650', 'GeForce GTX 1660') for execution_device in execution_devices: diff --git a/facefusion/system.py b/facefusion/system.py index 745f1aa4..61a8438b 100644 --- a/facefusion/system.py +++ b/facefusion/system.py @@ -7,7 +7,7 @@ from typing import List import psutil from facefusion import state_manager -from facefusion.types import DiskMetrics, ExecutionDevice, MemoryMetrics, Metrics, NetworkMetrics, ProcessorMetrics +from facefusion.types import DiskMetrics, ExecutionProvider, GraphicDevice, MemoryMetrics, Metrics, NetworkMetrics, ProcessorMetrics def get_metrics_set() -> Metrics: @@ -15,7 +15,7 @@ def get_metrics_set() -> Metrics: return\ { - 'execution_devices': detect_graphic_devices(), + 'graphic_devices': detect_graphic_devices(state_manager.get_item('execution_providers')), 'disks': detect_disk_metrics([ drive_path ]), 'memory': detect_memory_metrics(), 'network': detect_network_metrics(), @@ -24,164 +24,156 @@ def get_metrics_set() -> Metrics: @lru_cache() -def detect_static_graphic_devices() -> List[ExecutionDevice]: - return detect_graphic_devices() +def detect_static_graphic_devices(execution_providers : List[ExecutionProvider]) -> List[GraphicDevice]: + return detect_graphic_devices(execution_providers) -def detect_graphic_devices() -> List[ExecutionDevice]: - execution_providers = state_manager.get_item('execution_providers') - +def detect_graphic_devices(execution_providers : List[ExecutionProvider]) -> List[GraphicDevice]: if any(execution_provider in [ 'rocm', 'migraphx' ] for execution_provider in execution_providers): - return detect_amd_devices() + return detect_amd_graphic_devices() if any(execution_provider in [ 'cuda', 'tensorrt' ] for execution_provider in execution_providers): - return detect_nvidia_devices() + return detect_nvidia_graphic_devices() return [] -def detect_nvidia_devices() -> List[ExecutionDevice]: +def detect_nvidia_graphic_devices() -> List[GraphicDevice]: pynvml = importlib.import_module('pynvml') - nvidia_devices : List[ExecutionDevice] = [] + graphic_devices : List[GraphicDevice] = [] - try: - pynvml.nvmlInit() - device_count = pynvml.nvmlDeviceGetCount() + pynvml.nvmlInit() + device_count = pynvml.nvmlDeviceGetCount() - for device_id in range(device_count): - handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) + for device_id in range(device_count): + handle = pynvml.nvmlDeviceGetHandleByIndex(device_id) - nvidia_devices.append( + graphic_devices.append( + { + 'driver_version': pynvml.nvmlSystemGetDriverVersion(), + 'framework': { - 'driver_version': pynvml.nvmlSystemGetDriverVersion(), - 'framework': + 'name': 'CUDA', + 'version': pynvml.nvmlSystemGetCudaDriverVersion() + }, + 'product': + { + 'vendor': 'NVIDIA', + 'name': pynvml.nvmlDeviceGetName(handle) + }, + 'video_memory': + { + 'total': { - 'name': 'CUDA', - 'version': pynvml.nvmlSystemGetCudaDriverVersion() + 'value': pynvml.nvmlDeviceGetMemoryInfo(handle).total // (1024 * 1024 * 1024), + 'unit': 'GB' }, - 'product': + 'free': { - 'vendor': 'NVIDIA', - 'name': pynvml.nvmlDeviceGetName(handle) - }, - 'video_memory': - { - 'total': - { - 'value': pynvml.nvmlDeviceGetMemoryInfo(handle).total // (1024 * 1024 * 1024), - 'unit': 'GB' - }, - 'free': - { - 'value': pynvml.nvmlDeviceGetMemoryInfo(handle).free // (1024 * 1024 * 1024), - 'unit': 'GB' - } - }, - 'temperature': - { - 'gpu': - { - 'value': pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU), - 'unit': 'C' - }, - 'memory': - { - 'value': 0, - 'unit': '%' - } - }, - 'utilization': - { - 'gpu': - { - 'value': pynvml.nvmlDeviceGetUtilizationRates(handle).gpu, - 'unit': '%' - }, - 'memory': - { - 'value': pynvml.nvmlDeviceGetUtilizationRates(handle).memory, - 'unit': '%' - } + 'value': pynvml.nvmlDeviceGetMemoryInfo(handle).free // (1024 * 1024 * 1024), + 'unit': 'GB' } - }) + }, + 'temperature': + { + 'gpu': + { + 'value': pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU), + 'unit': 'C' + }, + 'memory': + { + 'value': 0, + 'unit': '%' + } + }, + 'utilization': + { + 'gpu': + { + 'value': pynvml.nvmlDeviceGetUtilizationRates(handle).gpu, + 'unit': '%' + }, + 'memory': + { + 'value': pynvml.nvmlDeviceGetUtilizationRates(handle).memory, + 'unit': '%' + } + } + }) - pynvml.nvmlShutdown() - except Exception: - pass + pynvml.nvmlShutdown() - return nvidia_devices + return graphic_devices -def detect_amd_devices() -> List[ExecutionDevice]: +def detect_amd_graphic_devices() -> List[GraphicDevice]: amdsmi = importlib.import_module('amdsmi') - amd_devices : List[ExecutionDevice] = [] + graphic_devices : List[GraphicDevice] = [] - try: - amdsmi.amdsmi_init() - handles = amdsmi.amdsmi_get_processor_handles() + amdsmi.amdsmi_init() + handles = amdsmi.amdsmi_get_processor_handles() - for handle in handles: - driver_info = amdsmi.amdsmi_get_gpu_driver_info(handle) - vram_usage = amdsmi.amdsmi_get_gpu_vram_usage(handle) - activity = amdsmi.amdsmi_get_gpu_activity(handle) + for handle in handles: + driver_info = amdsmi.amdsmi_get_gpu_driver_info(handle) + vram_usage = amdsmi.amdsmi_get_gpu_vram_usage(handle) + activity = amdsmi.amdsmi_get_gpu_activity(handle) - amd_devices.append( + graphic_devices.append( + { + 'driver_version': driver_info.get('driver_version', ''), + 'framework': { - 'driver_version': driver_info.get('driver_version', ''), - 'framework': + 'name': 'ROCm', + 'version': driver_info.get('driver_version', '') + }, + 'product': + { + 'vendor': 'AMD', + 'name': amdsmi.amdsmi_get_gpu_asic_info(handle).get('market_name', '') + }, + 'video_memory': + { + 'total': { - 'name': 'ROCm', - 'version': driver_info.get('driver_version', '') + 'value': vram_usage.get('vram_total', 0) // (1024 * 1024 * 1024), + 'unit': 'GB' }, - 'product': + 'free': { - 'vendor': 'AMD', - 'name': amdsmi.amdsmi_get_gpu_asic_info(handle).get('market_name', '') - }, - 'video_memory': - { - 'total': - { - 'value': vram_usage.get('vram_total', 0) // (1024 * 1024 * 1024), - 'unit': 'GB' - }, - 'free': - { - 'value': (vram_usage.get('vram_total', 0) - vram_usage.get('vram_used', 0)) // (1024 * 1024 * 1024), - 'unit': 'GB' - } - }, - 'temperature': - { - 'gpu': - { - 'value': amdsmi.amdsmi_get_temp_metric(handle, amdsmi.AmdSmiTemperatureType.EDGE, amdsmi.AmdSmiTemperatureMetric.CURRENT) // 1000, - 'unit': 'C' - }, - 'memory': - { - 'value': 0, - 'unit': '%' - } - }, - 'utilization': - { - 'gpu': - { - 'value': activity.get('gfx_activity', 0), - 'unit': '%' - }, - 'memory': - { - 'value': activity.get('umc_activity', 0), - 'unit': '%' - } + 'value': (vram_usage.get('vram_total', 0) - vram_usage.get('vram_used', 0)) // (1024 * 1024 * 1024), + 'unit': 'GB' } - }) + }, + 'temperature': + { + 'gpu': + { + 'value': amdsmi.amdsmi_get_temp_metric(handle, amdsmi.AmdSmiTemperatureType.EDGE, amdsmi.AmdSmiTemperatureMetric.CURRENT) // 1000, + 'unit': 'C' + }, + 'memory': + { + 'value': 0, + 'unit': '%' + } + }, + 'utilization': + { + 'gpu': + { + 'value': activity.get('gfx_activity', 0), + 'unit': '%' + }, + 'memory': + { + 'value': activity.get('umc_activity', 0), + 'unit': '%' + } + } + }) - amdsmi.amdsmi_shut_down() - except Exception: - pass + amdsmi.amdsmi_shut_down() - return amd_devices + return graphic_devices def detect_disk_metrics(drive_paths : List[str]) -> List[DiskMetrics]: diff --git a/facefusion/types.py b/facefusion/types.py index 2343b6e9..13f9ed04 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -289,7 +289,7 @@ ExecutionDeviceUtilization = TypedDict('ExecutionDeviceUtilization', 'gpu' : Optional[ValueAndUnit], 'memory' : Optional[ValueAndUnit] }) -ExecutionDevice = TypedDict('ExecutionDevice', +GraphicDevice = TypedDict('GraphicDevice', { 'driver_version' : str, 'framework' : ExecutionDeviceFramework, @@ -323,7 +323,7 @@ ProcessorMetrics = TypedDict('ProcessorMetrics', }) Metrics = TypedDict('Metrics', { - 'execution_devices' : List[ExecutionDevice], + 'graphic_devices' : List[GraphicDevice], 'disks' : List[DiskMetrics], 'memory' : MemoryMetrics, 'network' : NetworkMetrics, diff --git a/tests/test_api_metrics.py b/tests/test_api_metrics.py index f0d4dcf7..c2b830fa 100644 --- a/tests/test_api_metrics.py +++ b/tests/test_api_metrics.py @@ -167,9 +167,9 @@ def test_get_metrics(test_client : TestClient) -> None: assert metrics_response.status_code == 200 - assert metrics_body.get('execution_devices')[0].get('driver_version') == '555.42' - assert metrics_body.get('execution_devices')[0].get('product').get('name') == 'RTX 4090' - assert metrics_body.get('execution_devices')[0].get('video_memory').get('total').get('value') == 24 + assert metrics_body.get('graphic_devices')[0].get('driver_version') == '555.42' + assert metrics_body.get('graphic_devices')[0].get('product').get('name') == 'RTX 4090' + assert metrics_body.get('graphic_devices')[0].get('video_memory').get('total').get('value') == 24 assert metrics_body.get('disks')[0].get('total').get('value') == 500 assert metrics_body.get('disks')[0].get('free').get('unit') == 'GB' @@ -201,9 +201,9 @@ def test_websocket_metrics(test_client : TestClient) -> None: ]) as websocket: metrics_set = websocket.receive_json() - assert metrics_set.get('execution_devices')[0].get('driver_version') == '555.42' - assert metrics_set.get('execution_devices')[0].get('product').get('name') == 'RTX 4090' - assert metrics_set.get('execution_devices')[0].get('video_memory').get('total').get('value') == 24 + assert metrics_set.get('graphic_devices')[0].get('driver_version') == '555.42' + assert metrics_set.get('graphic_devices')[0].get('product').get('name') == 'RTX 4090' + assert metrics_set.get('graphic_devices')[0].get('video_memory').get('total').get('value') == 24 assert metrics_set.get('disks')[0].get('total').get('value') == 500 assert metrics_set.get('disks')[0].get('free').get('unit') == 'GB'