rename video memory to memory, add amdsmi library

facefusion/libraries/amd_smi.py

@@ -0,0 +1,146 @@
+import ctypes
+from functools import lru_cache
+from typing import List, Optional
+
+from facefusion.common_helper import is_linux
+
+
+def resolve_library_file() -> Optional[str]:
+	if is_linux():
+		return 'libamd_smi.so'
+	return None
+
+
+@lru_cache
+def create_static_library() -> Optional[ctypes.CDLL]:
+	library_file = resolve_library_file()
+
+	if library_file:
+		amd_smi_library = ctypes.CDLL(library_file)
+		return init_ctypes(amd_smi_library)
+
+	return None
+
+
+def create_asic_info_configuration() -> ctypes.Structure:
+	return type('AMDSMI_ASIC_INFO', (ctypes.Structure,),
+	{
+		'_pack_': 1,
+		'_fields_':
+		[
+			('market_name', ctypes.c_char * 256),
+			('vendor_id', ctypes.c_uint32),
+			('vendor_name', ctypes.c_char * 256),
+			('subvendor_id', ctypes.c_uint32),
+			('device_id', ctypes.c_uint64),
+			('rev_id', ctypes.c_uint32),
+			('asic_serial', ctypes.c_char * 256),
+			('oam_id', ctypes.c_uint32),
+			('num_of_compute_units', ctypes.c_uint32),
+			('padding', ctypes.c_ubyte * 4),
+			('target_graphics_version', ctypes.c_uint64),
+			('subsystem_id', ctypes.c_uint32),
+			('reserved', ctypes.c_uint32 * 21)
+		]
+	})()
+
+
+def create_driver_info_configuration() -> ctypes.Structure:
+	return type('AMDSMI_DRIVER_INFO', (ctypes.Structure,),
+	{
+		'_pack_': 1,
+		'_fields_':
+		[
+			('driver_version', ctypes.c_char * 256),
+			('driver_date', ctypes.c_char * 256),
+			('driver_name', ctypes.c_char * 256)
+		]
+	})()
+
+
+def create_memory_configuration() -> ctypes.Structure:
+	return type('AMDSMI_VRAM_USAGE', (ctypes.Structure,),
+	{
+		'_pack_': 1,
+		'_fields_':
+		[
+			('vram_total', ctypes.c_uint32),
+			('vram_used', ctypes.c_uint32),
+			('reserved', ctypes.c_uint32 * 2)
+		]
+	})()
+
+
+def create_utilization_configuration() -> ctypes.Structure:
+	return type('AMDSMI_ENGINE_USAGE', (ctypes.Structure,),
+	{
+		'_pack_': 1,
+		'_fields_':
+		[
+			('gfx_activity', ctypes.c_uint32),
+			('umc_activity', ctypes.c_uint32),
+			('mm_activity', ctypes.c_uint32),
+			('reserved', ctypes.c_uint32 * 13)
+		]
+	})()
+
+
+def init_ctypes(amd_smi : ctypes.CDLL) -> ctypes.CDLL:
+	void_pointer = ctypes.POINTER(None)
+
+	amd_smi.amdsmi_init.argtypes = [ ctypes.c_uint64 ]
+	amd_smi.amdsmi_init.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_shut_down.argtypes = []
+	amd_smi.amdsmi_shut_down.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_socket_handles.argtypes = [ ctypes.POINTER(ctypes.c_uint32), ctypes.POINTER(void_pointer) ]
+	amd_smi.amdsmi_get_socket_handles.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_processor_handles.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint32), ctypes.POINTER(void_pointer) ]
+	amd_smi.amdsmi_get_processor_handles.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_gpu_driver_info.argtypes = [ ctypes.c_void_p, ctypes.c_void_p ]
+	amd_smi.amdsmi_get_gpu_driver_info.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_gpu_vram_usage.argtypes = [ ctypes.c_void_p, ctypes.c_void_p ]
+	amd_smi.amdsmi_get_gpu_vram_usage.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_gpu_memory_total.argtypes = [ ctypes.c_void_p, ctypes.c_uint32, ctypes.POINTER(ctypes.c_uint64) ]
+	amd_smi.amdsmi_get_gpu_memory_total.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_gpu_memory_usage.argtypes = [ ctypes.c_void_p, ctypes.c_uint32, ctypes.POINTER(ctypes.c_uint64) ]
+	amd_smi.amdsmi_get_gpu_memory_usage.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_gpu_activity.argtypes = [ ctypes.c_void_p, ctypes.c_void_p ]
+	amd_smi.amdsmi_get_gpu_activity.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_gpu_asic_info.argtypes = [ ctypes.c_void_p, ctypes.c_void_p ]
+	amd_smi.amdsmi_get_gpu_asic_info.restype = ctypes.c_uint32
+
+	amd_smi.amdsmi_get_temp_metric.argtypes = [ ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.POINTER(ctypes.c_int64) ]
+	amd_smi.amdsmi_get_temp_metric.restype = ctypes.c_uint32
+
+	return amd_smi
+
+
+def find_device_handles(amd_smi_library : ctypes.CDLL) -> List[ctypes.c_void_p]:
+	device_handles : List[ctypes.c_void_p] = []
+
+	socket_count = ctypes.c_uint32(0)
+	amd_smi_library.amdsmi_get_socket_handles(ctypes.byref(socket_count), ctypes.POINTER(ctypes.c_void_p)())
+	socket_handles = (ctypes.c_void_p * socket_count.value)()
+	amd_smi_library.amdsmi_get_socket_handles(ctypes.byref(socket_count), socket_handles)
+
+	for socket_index in range(socket_count.value):
+		device_count = ctypes.c_uint32(0)
+		amd_smi_library.amdsmi_get_processor_handles(socket_handles[socket_index], ctypes.byref(device_count), ctypes.POINTER(ctypes.c_void_p)())
+		processor_handles = (ctypes.c_void_p * device_count.value)()
+		amd_smi_library.amdsmi_get_processor_handles(socket_handles[socket_index], ctypes.byref(device_count), processor_handles)
+
+		for device_index in range(device_count.value):
+			device_handles.append(ctypes.c_void_p(processor_handles[device_index]))
+
+	return device_handles
+
+

facefusion/libraries/nvidia_ml.py

@@ -1,6 +1,6 @@
 import ctypes
 from functools import lru_cache
-from typing import Optional
+from typing import List, Optional
 
 from facefusion.common_helper import is_linux, is_windows
 
@@ -79,3 +79,19 @@ def init_ctypes(nvidia_ml : ctypes.CDLL) -> ctypes.CDLL:
 	nvidia_ml.nvmlDeviceGetUtilizationRates.restype = ctypes.c_int
 
 	return nvidia_ml
+
+
+def find_device_handles(nvidia_ml_library : ctypes.CDLL) -> List[ctypes.c_void_p]:
+	device_handles : List[ctypes.c_void_p] = []
+
+	device_count = ctypes.c_uint()
+	nvidia_ml_library.nvmlDeviceGetCount_v2(ctypes.byref(device_count))
+
+	for device_id in range(device_count.value):
+		device_handle = ctypes.c_void_p()
+		nvidia_ml_library.nvmlDeviceGetHandleByIndex_v2(device_id, ctypes.byref(device_handle))
+		device_handles.append(device_handle)
+
+	return device_handles
+
+

facefusion/system.py

@@ -1,5 +1,4 @@
 import ctypes
-import importlib
 import shutil
 from pathlib import Path
 from typing import List
@@ -7,7 +6,7 @@ from typing import List
 import psutil
 
 from facefusion import state_manager
-from facefusion.libraries import nvidia_ml as nvidia_ml_module
+from facefusion.libraries import amd_smi as amd_smi_module, nvidia_ml as nvidia_ml_module
 from facefusion.types import DiskMetrics, ExecutionProvider, GraphicDevice, MemoryMetrics, Metrics, NetworkMetrics, ProcessorMetrics
 
 
@@ -39,19 +38,13 @@ def detect_nvidia_graphic_devices() -> List[GraphicDevice]:
 	if nvidia_ml_library:
 		nvidia_ml_library.nvmlInit_v2()
 
-		device_count = ctypes.c_uint()
-		nvidia_ml_library.nvmlDeviceGetCount_v2(ctypes.byref(device_count))
-
 		driver_version = ctypes.create_string_buffer(80)
 		nvidia_ml_library.nvmlSystemGetDriverVersion(driver_version, 80)
 
 		cuda_version = ctypes.c_int()
 		nvidia_ml_library.nvmlSystemGetCudaDriverVersion(ctypes.byref(cuda_version))
 
-		for device_id in range(device_count.value):
-			device_handle = ctypes.c_void_p()
-			nvidia_ml_library.nvmlDeviceGetHandleByIndex_v2(device_id, ctypes.byref(device_handle))
-
+		for device_handle in nvidia_ml_module.find_device_handles(nvidia_ml_library):
 			name = ctypes.create_string_buffer(96)
 			nvidia_ml_library.nvmlDeviceGetName(device_handle, name, 96)
 
@@ -77,7 +70,7 @@ def detect_nvidia_graphic_devices() -> List[GraphicDevice]:
 					'vendor': 'NVIDIA',
 					'name': name.value.decode()
 				},
-				'video_memory':
+				'memory':
 				{
 					'total':
 					{
@@ -124,72 +117,86 @@ def detect_nvidia_graphic_devices() -> List[GraphicDevice]:
 
 
 def detect_amd_graphic_devices() -> List[GraphicDevice]:
-	amdsmi = importlib.import_module('amdsmi')
+	amd_smi_library = amd_smi_module.create_static_library()
 	graphic_devices : List[GraphicDevice] = []
 
-	amdsmi.amdsmi_init()
-	handles = amdsmi.amdsmi_get_processor_handles()
+	if amd_smi_library:
+		amd_smi_library.amdsmi_init(ctypes.c_uint64(2))
 
-	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 device_handle in amd_smi_module.find_device_handles(amd_smi_library):
+			driver_info = amd_smi_module.create_driver_info_configuration()
+			amd_smi_library.amdsmi_get_gpu_driver_info(device_handle, ctypes.byref(driver_info))
 
-		graphic_devices.append(
-		{
-			'driver_version': driver_info.get('driver_version', ''),
-			'framework':
+			asic_info = amd_smi_module.create_asic_info_configuration()
+			amd_smi_library.amdsmi_get_gpu_asic_info(device_handle, ctypes.byref(asic_info))
+
+			memory_total = ctypes.c_uint64()
+			amd_smi_library.amdsmi_get_gpu_memory_total(device_handle, 0, ctypes.byref(memory_total))
+
+			memory_usage = ctypes.c_uint64()
+			amd_smi_library.amdsmi_get_gpu_memory_usage(device_handle, 0, ctypes.byref(memory_usage))
+
+			temperature = ctypes.c_int64()
+			amd_smi_library.amdsmi_get_temp_metric(device_handle, 0, 0, ctypes.byref(temperature))
+
+			utilization = amd_smi_module.create_utilization_configuration()
+			amd_smi_library.amdsmi_get_gpu_activity(device_handle, ctypes.byref(utilization))
+
+			graphic_devices.append(
 			{
-				'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':
+				'driver_version': driver_info.driver_version.decode(),
+				'framework':
 				{
-					'value': vram_usage.get('vram_total', 0) // (1024 * 1024 * 1024),
-					'unit': 'GB'
+					'name': 'ROCm',
+					'version': driver_info.driver_version.decode()
 				},
-				'used':
+				'product':
 				{
-					'value': 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'
+					'vendor': 'AMD',
+					'name': asic_info.market_name.decode()
 				},
 				'memory':
 				{
-					'value': 0,
-					'unit': '%'
-				}
-			},
-			'utilization':
-			{
-				'gpu':
-				{
-					'value': activity.get('gfx_activity', 0),
-					'unit': '%'
+					'total':
+					{
+						'value': memory_total.value // (1024 * 1024 * 1024),
+						'unit': 'GB'
+					},
+					'used':
+					{
+						'value': memory_usage.value // (1024 * 1024 * 1024),
+						'unit': 'GB'
+					}
 				},
-				'memory':
+				'temperature':
 				{
-					'value': activity.get('umc_activity', 0),
-					'unit': '%'
+					'gpu':
+					{
+						'value': temperature.value // 1000,
+						'unit': 'C'
+					},
+					'memory':
+					{
+						'value': 0,
+						'unit': '%'
+					}
+				},
+				'utilization':
+				{
+					'gpu':
+					{
+						'value': utilization.gfx_activity,
+						'unit': '%'
+					},
+					'memory':
+					{
+						'value': utilization.umc_activity,
+						'unit': '%'
+					}
 				}
-			}
-		})
+			})
 
-	amdsmi.amdsmi_shut_down()
+		amd_smi_library.amdsmi_shut_down()
 
 	return graphic_devices
 

facefusion/types.py

@@ -311,7 +311,7 @@ ExecutionDeviceProduct = TypedDict('ExecutionDeviceProduct',
 	'vendor' : str,
 	'name' : str
 })
-ExecutionDeviceVideoMemory = TypedDict('ExecutionDeviceVideoMemory',
+ExecutionDeviceMemory = TypedDict('ExecutionDeviceMemory',
 {
 	'total' : Optional[ValueAndUnit],
 	'used' : Optional[ValueAndUnit]
@@ -331,7 +331,7 @@ GraphicDevice = TypedDict('GraphicDevice',
 	'driver_version' : str,
 	'framework' : ExecutionDeviceFramework,
 	'product' : ExecutionDeviceProduct,
-	'video_memory' : ExecutionDeviceVideoMemory,
+	'memory' : ExecutionDeviceMemory,
 	'temperature' : ExecutionDeviceTemperature,
 	'utilization' : ExecutionDeviceUtilization
 })

requirements.txt

@@ -1,4 +1,3 @@
-amdsmi==7.0.2
 numpy==2.4.2
 onnx==1.20.1
 onnxruntime==1.24.1

tests/test_api_metrics.py

@@ -105,7 +105,7 @@ def mock_detect_execution_devices(mocker : MockerFixture) -> None:
 				'vendor': 'NVIDIA',
 				'name': 'RTX 4090'
 			},
-			'video_memory':
+			'memory':
 			{
 				'total':
 				{
@@ -169,8 +169,8 @@ def test_get_metrics(test_client : TestClient) -> None:
 
 	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('graphic_devices')[0].get('video_memory').get('used').get('value') == 4
+	assert metrics_body.get('graphic_devices')[0].get('memory').get('total').get('value') == 24
+	assert metrics_body.get('graphic_devices')[0].get('memory').get('used').get('value') == 4
 
 	assert metrics_body.get('disks')[0].get('total').get('value') == 500
 	assert metrics_body.get('disks')[0].get('free').get('unit') == 'GB'
@@ -221,8 +221,8 @@ def test_websocket_metrics(test_client : TestClient) -> None:
 
 		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('graphic_devices')[0].get('video_memory').get('used').get('value') == 4
+		assert metrics_set.get('graphic_devices')[0].get('memory').get('total').get('value') == 24
+		assert metrics_set.get('graphic_devices')[0].get('memory').get('used').get('value') == 4
 
 		assert metrics_set.get('disks')[0].get('total').get('value') == 500
 		assert metrics_set.get('disks')[0].get('free').get('unit') == 'GB'