Spaces:
Paused
Paused
| from typing import Any, List | |
| import onnxruntime | |
| def encode_execution_providers(execution_providers : List[str]) -> List[str]: | |
| return [ execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers ] | |
| def decode_execution_providers(execution_providers: List[str]) -> List[str]: | |
| available_execution_providers = onnxruntime.get_available_providers() | |
| encoded_execution_providers = encode_execution_providers(available_execution_providers) | |
| return [ execution_provider for execution_provider, encoded_execution_provider in zip(available_execution_providers, encoded_execution_providers) if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers) ] | |
| def apply_execution_provider_options(execution_providers: List[str]) -> List[Any]: | |
| execution_providers_with_options : List[Any] = [] | |
| for execution_provider in execution_providers: | |
| if execution_provider == 'CUDAExecutionProvider': | |
| execution_providers_with_options.append((execution_provider, | |
| { | |
| 'cudnn_conv_algo_search': 'DEFAULT' | |
| })) | |
| else: | |
| execution_providers_with_options.append(execution_provider) | |
| return execution_providers_with_options | |
| def map_torch_backend(execution_providers : List[str]) -> str: | |
| if 'CoreMLExecutionProvider' in execution_providers: | |
| return 'mps' | |
| if 'CUDAExecutionProvider' in execution_providers or 'ROCMExecutionProvider' in execution_providers : | |
| return 'cuda' | |
| if 'OpenVINOExecutionProvider' in execution_providers: | |
| return 'mkl' | |
| return 'cpu' | |