Spaces:
Build error
Build error
File size: 4,849 Bytes
a1da63c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import subprocess
import xml.etree.ElementTree as ElementTree
from functools import lru_cache
from typing import Any, List
from onnxruntime import get_available_providers, set_default_logger_severity
from facefusion.choices import execution_provider_set
from facefusion.typing import ExecutionDevice, ExecutionProviderKey, ExecutionProviderSet, ExecutionProviderValue, ValueAndUnit
set_default_logger_severity(3)
def get_execution_provider_choices() -> List[ExecutionProviderKey]:
return list(get_available_execution_provider_set().keys())
def has_execution_provider(execution_provider_key : ExecutionProviderKey) -> bool:
return execution_provider_key in get_execution_provider_choices()
def get_available_execution_provider_set() -> ExecutionProviderSet:
available_execution_providers = get_available_providers()
available_execution_provider_set : ExecutionProviderSet = {}
for execution_provider_key, execution_provider_value in execution_provider_set.items():
if execution_provider_value in available_execution_providers:
available_execution_provider_set[execution_provider_key] = execution_provider_value
return available_execution_provider_set
def extract_execution_providers(execution_provider_keys : List[ExecutionProviderKey]) -> List[ExecutionProviderValue]:
return [ execution_provider_set[execution_provider_key] for execution_provider_key in execution_provider_keys if execution_provider_key in execution_provider_set ]
def create_execution_providers(execution_device_id : str, execution_provider_keys : List[ExecutionProviderKey]) -> List[Any]:
execution_providers = extract_execution_providers(execution_provider_keys)
execution_providers_with_options : List[Any] = []
for execution_provider in execution_providers:
if execution_provider == 'CUDAExecutionProvider':
execution_providers_with_options.append((execution_provider,
{
'device_id': execution_device_id,
'cudnn_conv_algo_search': 'EXHAUSTIVE' if use_exhaustive() else 'DEFAULT'
}))
elif execution_provider == 'TensorrtExecutionProvider':
execution_providers_with_options.append((execution_provider,
{
'device_id': execution_device_id,
'trt_engine_cache_enable': True,
'trt_engine_cache_path': '.caches',
'trt_timing_cache_enable': True,
'trt_timing_cache_path': '.caches'
}))
elif execution_provider == 'OpenVINOExecutionProvider':
execution_providers_with_options.append((execution_provider,
{
'device_type': 'GPU.' + execution_device_id,
'precision': 'FP32'
}))
elif execution_provider in [ 'DmlExecutionProvider', 'ROCMExecutionProvider' ]:
execution_providers_with_options.append((execution_provider,
{
'device_id': execution_device_id
}))
elif execution_provider == 'CoreMLExecutionProvider':
execution_providers_with_options.append(execution_provider)
if 'CPUExecutionProvider' in execution_providers:
execution_providers_with_options.append('CPUExecutionProvider')
return execution_providers_with_options
def use_exhaustive() -> bool:
execution_devices = detect_static_execution_devices()
product_names = ('GeForce GTX 1630', 'GeForce GTX 1650', 'GeForce GTX 1660')
return any(execution_device.get('product').get('name').startswith(product_names) for execution_device in execution_devices)
def run_nvidia_smi() -> subprocess.Popen[bytes]:
commands = [ 'nvidia-smi', '--query', '--xml-format' ]
return subprocess.Popen(commands, stdout = subprocess.PIPE)
@lru_cache(maxsize = None)
def detect_static_execution_devices() -> List[ExecutionDevice]:
return detect_execution_devices()
def detect_execution_devices() -> List[ExecutionDevice]:
execution_devices : List[ExecutionDevice] = []
try:
output, _ = run_nvidia_smi().communicate()
root_element = ElementTree.fromstring(output)
except Exception:
root_element = ElementTree.Element('xml')
for gpu_element in root_element.findall('gpu'):
execution_devices.append(
{
'driver_version': root_element.find('driver_version').text,
'framework':
{
'name': 'CUDA',
'version': root_element.find('cuda_version').text
},
'product':
{
'vendor': 'NVIDIA',
'name': gpu_element.find('product_name').text.replace('NVIDIA ', '')
},
'video_memory':
{
'total': create_value_and_unit(gpu_element.find('fb_memory_usage/total').text),
'free': create_value_and_unit(gpu_element.find('fb_memory_usage/free').text)
},
'utilization':
{
'gpu': create_value_and_unit(gpu_element.find('utilization/gpu_util').text),
'memory': create_value_and_unit(gpu_element.find('utilization/memory_util').text)
}
})
return execution_devices
def create_value_and_unit(text : str) -> ValueAndUnit:
value, unit = text.split()
value_and_unit : ValueAndUnit =\
{
'value': int(value),
'unit': str(unit)
}
return value_and_unit
|