add memory metrics

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': '%'
+		}
+	}

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']

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