diff --git a/facefusion/system.py b/facefusion/system.py index 9b3ba33f..608d63f6 100644 --- a/facefusion/system.py +++ b/facefusion/system.py @@ -2,9 +2,11 @@ import shutil from pathlib import Path from typing import List +import psutil + from facefusion import state_manager from facefusion.execution import detect_execution_devices -from facefusion.types import DiskMetrics, Metrics +from facefusion.types import DiskMetrics, MemoryMetrics, Metrics def get_metrics_set() -> Metrics: @@ -13,7 +15,8 @@ def get_metrics_set() -> Metrics: return\ { 'execution_devices': detect_execution_devices(), - 'disks': detect_disk_metrics([ drive_path ]) + 'disks': detect_disk_metrics([ drive_path ]), + 'memory': detect_memory_metrics() } @@ -43,3 +46,26 @@ def detect_disk_metrics(drive_paths : List[str]) -> List[DiskMetrics]: }) return disk_metrics + + +def detect_memory_metrics() -> MemoryMetrics: + virtual_memory = psutil.virtual_memory() + + return\ + { + 'total': + { + 'value': int(virtual_memory.total / (1024 * 1024 * 1024)), + 'unit': 'GB' + }, + 'free': + { + 'value': int(virtual_memory.available / (1024 * 1024 * 1024)), + 'unit': 'GB' + }, + 'utilization': + { + 'value': int(virtual_memory.percent), + 'unit': '%' + } + } diff --git a/facefusion/types.py b/facefusion/types.py index bad1a3c3..e2a19bdc 100755 --- a/facefusion/types.py +++ b/facefusion/types.py @@ -304,10 +304,17 @@ DiskMetrics = TypedDict('DiskMetrics', 'free' : ValueAndUnit, 'utilization' : ValueAndUnit }) +MemoryMetrics = TypedDict('MemoryMetrics', +{ + 'total' : ValueAndUnit, + 'free' : ValueAndUnit, + 'utilization' : ValueAndUnit +}) Metrics = TypedDict('Metrics', { 'execution_devices' : List[ExecutionDevice], - 'disks' : List[DiskMetrics] + 'disks' : List[DiskMetrics], + 'memory' : MemoryMetrics }) DownloadProvider = Literal['github', 'huggingface'] diff --git a/tests/test_api_metrics.py b/tests/test_api_metrics.py index 6bcf2e14..0165e502 100644 --- a/tests/test_api_metrics.py +++ b/tests/test_api_metrics.py @@ -42,6 +42,24 @@ def mock_detect_execution_devices(mocker : MockerFixture) -> None: } } ]) + mocker.patch('facefusion.system.detect_memory_metrics', return_value = + { + 'total': + { + 'value': 32, + 'unit': 'GB' + }, + 'free': + { + 'value': 16, + 'unit': 'GB' + }, + 'utilization': + { + 'value': 50, + 'unit': '%' + } + }) mocker.patch('facefusion.system.detect_execution_devices', return_value = [ { @@ -126,6 +144,10 @@ def test_get_metrics(test_client : TestClient) -> None: assert metrics_body.get('disks')[0].get('free').get('unit') == 'GB' assert metrics_body.get('disks')[0].get('utilization').get('value') == 60 + assert metrics_body.get('memory').get('total').get('value') == 32 + assert metrics_body.get('memory').get('free').get('unit') == 'GiB' + assert metrics_body.get('memory').get('utilization').get('value') == 50 + def test_websocket_metrics(test_client : TestClient) -> None: create_session_response = test_client.post('/session', json = @@ -147,3 +169,7 @@ def test_websocket_metrics(test_client : TestClient) -> None: assert metrics_set.get('disks')[0].get('total').get('value') == 500 assert metrics_set.get('disks')[0].get('free').get('unit') == 'GB' assert metrics_set.get('disks')[0].get('utilization').get('value') == 60 + + assert metrics_set.get('memory').get('total').get('value') == 32 + assert metrics_set.get('memory').get('free').get('unit') == 'GiB' + assert metrics_set.get('memory').get('utilization').get('value') == 50