This commit is contained in:
henryruhs
2026-05-08 10:44:27 +02:00
parent 949d9cd276
commit b549a92a35
2 changed files with 164 additions and 59 deletions
+81
View File
@@ -0,0 +1,81 @@
import ctypes
from functools import lru_cache
from typing import Optional
from facefusion.common_helper import is_linux, is_windows
def resolve_library_file() -> Optional[str]:
if is_linux():
return 'libnvidia-ml.so.1'
if is_windows():
return 'nvml.dll'
return None
@lru_cache
def create_static_library() -> Optional[ctypes.CDLL]:
library_file = resolve_library_file()
if library_file:
nvml_library = ctypes.CDLL(library_file)
return init_ctypes(nvml_library)
return None
def create_memory_configuration() -> ctypes.Structure:
return type('NVML_MEMORY', (ctypes.Structure,),
{
'_fields_':
[
('total', ctypes.c_ulonglong),
('free', ctypes.c_ulonglong),
('used', ctypes.c_ulonglong)
]
})()
def create_utilization_configuration() -> ctypes.Structure:
return type('NVML_UTILIZATION', (ctypes.Structure,),
{
'_fields_':
[
('gpu', ctypes.c_uint),
('memory', ctypes.c_uint)
]
})()
def init_ctypes(nvidia_ml : ctypes.CDLL) -> ctypes.CDLL:
nvidia_ml.nvmlInit_v2.argtypes = []
nvidia_ml.nvmlInit_v2.restype = ctypes.c_int
nvidia_ml.nvmlShutdown.argtypes = []
nvidia_ml.nvmlShutdown.restype = ctypes.c_int
nvidia_ml.nvmlDeviceGetCount_v2.argtypes = [ctypes.POINTER(ctypes.c_uint)]
nvidia_ml.nvmlDeviceGetCount_v2.restype = ctypes.c_int
nvidia_ml.nvmlSystemGetDriverVersion.argtypes = [ctypes.c_char_p, ctypes.c_uint]
nvidia_ml.nvmlSystemGetDriverVersion.restype = ctypes.c_int
nvidia_ml.nvmlSystemGetCudaDriverVersion.argtypes = [ctypes.POINTER(ctypes.c_int)]
nvidia_ml.nvmlSystemGetCudaDriverVersion.restype = ctypes.c_int
nvidia_ml.nvmlDeviceGetHandleByIndex_v2.argtypes = [ctypes.c_uint, ctypes.POINTER(ctypes.c_void_p)]
nvidia_ml.nvmlDeviceGetHandleByIndex_v2.restype = ctypes.c_int
nvidia_ml.nvmlDeviceGetName.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint]
nvidia_ml.nvmlDeviceGetName.restype = ctypes.c_int
nvidia_ml.nvmlDeviceGetMemoryInfo.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
nvidia_ml.nvmlDeviceGetMemoryInfo.restype = ctypes.c_int
nvidia_ml.nvmlDeviceGetTemperature.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.POINTER(ctypes.c_uint)]
nvidia_ml.nvmlDeviceGetTemperature.restype = ctypes.c_int
nvidia_ml.nvmlDeviceGetUtilizationRates.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
nvidia_ml.nvmlDeviceGetUtilizationRates.restype = ctypes.c_int
return nvidia_ml
+83 -59
View File
@@ -1,3 +1,4 @@
import ctypes
import importlib
import shutil
from pathlib import Path
@@ -6,6 +7,7 @@ from typing import List
import psutil
from facefusion import state_manager
from facefusion.libraries.nvidia_ml import create_memory_configuration, create_utilization_configuration, create_static_library
from facefusion.types import DiskMetrics, ExecutionProvider, GraphicDevice, MemoryMetrics, Metrics, NetworkMetrics, ProcessorMetrics
@@ -31,70 +33,92 @@ def detect_graphic_devices(execution_providers : List[ExecutionProvider]) -> Lis
def detect_nvidia_graphic_devices() -> List[GraphicDevice]:
pynvml = importlib.import_module('pynvml')
nvidia_ml = create_static_library()
graphic_devices : List[GraphicDevice] = []
pynvml.nvmlInit()
device_count = pynvml.nvmlDeviceGetCount()
if nvidia_ml:
nvidia_ml.nvmlInit_v2()
for device_id in range(device_count):
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
device_count = ctypes.c_uint()
nvidia_ml.nvmlDeviceGetCount_v2(ctypes.byref(device_count))
graphic_devices.append(
{
'driver_version': pynvml.nvmlSystemGetDriverVersion(),
'framework':
{
'name': 'CUDA',
'version': pynvml.nvmlSystemGetCudaDriverVersion()
},
'product':
{
'vendor': 'NVIDIA',
'name': pynvml.nvmlDeviceGetName(handle)
},
'video_memory':
{
'total':
{
'value': pynvml.nvmlDeviceGetMemoryInfo(handle).total // (1024 * 1024 * 1024),
'unit': 'GB'
},
'used':
{
'value': pynvml.nvmlDeviceGetMemoryInfo(handle).used // (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': '%'
}
}
})
driver_version = ctypes.create_string_buffer(80)
nvidia_ml.nvmlSystemGetDriverVersion(driver_version, 80)
pynvml.nvmlShutdown()
cuda_version = ctypes.c_int()
nvidia_ml.nvmlSystemGetCudaDriverVersion(ctypes.byref(cuda_version))
for device_id in range(device_count.value):
device_handle = ctypes.c_void_p()
nvidia_ml.nvmlDeviceGetHandleByIndex_v2(device_id, ctypes.byref(device_handle))
name = ctypes.create_string_buffer(96)
nvidia_ml.nvmlDeviceGetName(device_handle, name, 96)
memory = create_memory_configuration()
nvidia_ml.nvmlDeviceGetMemoryInfo(device_handle, ctypes.byref(memory))
temperature = ctypes.c_uint()
nvidia_ml.nvmlDeviceGetTemperature(device_handle, 0, ctypes.byref(temperature))
utilization = create_utilization_configuration()
nvidia_ml.nvmlDeviceGetUtilizationRates(device_handle, ctypes.byref(utilization))
graphic_devices.append(
{
'driver_version': driver_version.value.decode(),
'framework':
{
'name': 'CUDA',
'version': cuda_version.value
},
'product':
{
'vendor': 'NVIDIA',
'name': name.value.decode()
},
'video_memory':
{
'total':
{
'value': memory.total // (1024 * 1024 * 1024),
'unit': 'GB'
},
'used':
{
'value': memory.used // (1024 * 1024 * 1024),
'unit': 'GB'
}
},
'temperature':
{
'gpu':
{
'value': temperature.value,
'unit': 'C'
},
'memory':
{
'value': 0,
'unit': '%'
}
},
'utilization':
{
'gpu':
{
'value': utilization.gpu,
'unit': '%'
},
'memory':
{
'value': utilization.memory,
'unit': '%'
}
}
})
nvidia_ml.nvmlShutdown()
return graphic_devices