### What problem does this PR solve? Added cuda_is_available ### Type of change - [x] Refactoringtags/v0.16.0
| import math | import math | ||||
| import numpy as np | import numpy as np | ||||
| import cv2 | import cv2 | ||||
| import torch | |||||
| import onnxruntime as ort | import onnxruntime as ort | ||||
| from .postprocess import build_post_process | from .postprocess import build_post_process | ||||
| raise ValueError("not find model file path {}".format( | raise ValueError("not find model file path {}".format( | ||||
| model_file_path)) | model_file_path)) | ||||
| def cuda_is_available(): | |||||
| try: | |||||
| import torch | |||||
| if torch.cuda.is_available(): | |||||
| return True | |||||
| except Exception: | |||||
| return False | |||||
| return False | |||||
| options = ort.SessionOptions() | options = ort.SessionOptions() | ||||
| options.enable_cpu_mem_arena = False | options.enable_cpu_mem_arena = False | ||||
| options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL | options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL | ||||
| # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580 | # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580 | ||||
| # Shrink GPU memory after execution | # Shrink GPU memory after execution | ||||
| run_options = ort.RunOptions() | run_options = ort.RunOptions() | ||||
| if torch.cuda.is_available(): | |||||
| if cuda_is_available(): | |||||
| cuda_provider_options = { | cuda_provider_options = { | ||||
| "device_id": 0, # Use specific GPU | "device_id": 0, # Use specific GPU | ||||
| "gpu_mem_limit": 512 * 1024 * 1024, # Limit gpu memory | "gpu_mem_limit": 512 * 1024 * 1024, # Limit gpu memory |
| import cv2 | import cv2 | ||||
| from copy import deepcopy | from copy import deepcopy | ||||
| import torch | |||||
| import onnxruntime as ort | import onnxruntime as ort | ||||
| from huggingface_hub import snapshot_download | from huggingface_hub import snapshot_download | ||||
| if not os.path.exists(model_file_path): | if not os.path.exists(model_file_path): | ||||
| raise ValueError("not find model file path {}".format( | raise ValueError("not find model file path {}".format( | ||||
| model_file_path)) | model_file_path)) | ||||
| def cuda_is_available(): | |||||
| try: | |||||
| import torch | |||||
| if torch.cuda.is_available(): | |||||
| return True | |||||
| except Exception: | |||||
| return False | |||||
| return False | |||||
| # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580 | # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580 | ||||
| # Shrink GPU memory after execution | # Shrink GPU memory after execution | ||||
| self.run_options = ort.RunOptions() | self.run_options = ort.RunOptions() | ||||
| if torch.cuda.is_available(): | |||||
| if cuda_is_available(): | |||||
| options = ort.SessionOptions() | options = ort.SessionOptions() | ||||
| options.enable_cpu_mem_arena = False | options.enable_cpu_mem_arena = False | ||||
| cuda_provider_options = { | cuda_provider_options = { |
| "tencentcloud-sdk-python==3.0.1215", | "tencentcloud-sdk-python==3.0.1215", | ||||
| "tika==2.6.0", | "tika==2.6.0", | ||||
| "tiktoken==0.7.0", | "tiktoken==0.7.0", | ||||
| "torch>=2.5.0,<3.0.0", | |||||
| "umap_learn==0.5.6", | "umap_learn==0.5.6", | ||||
| "vertexai==1.64.0", | "vertexai==1.64.0", | ||||
| "volcengine==1.0.146", | "volcengine==1.0.146", | ||||
| "fastembed>=0.3.6,<0.4.0; sys_platform == 'darwin' or platform_machine != 'x86_64'", | "fastembed>=0.3.6,<0.4.0; sys_platform == 'darwin' or platform_machine != 'x86_64'", | ||||
| "fastembed-gpu>=0.3.6,<0.4.0; sys_platform != 'darwin' and platform_machine == 'x86_64'", | "fastembed-gpu>=0.3.6,<0.4.0; sys_platform != 'darwin' and platform_machine == 'x86_64'", | ||||
| "flagembedding==1.2.10", | "flagembedding==1.2.10", | ||||
| "torch>=2.5.0,<3.0.0", | |||||
| "transformers>=4.35.0,<5.0.0" | "transformers>=4.35.0,<5.0.0" | ||||
| ] | ] |
| { name = "tencentcloud-sdk-python" }, | { name = "tencentcloud-sdk-python" }, | ||||
| { name = "tika" }, | { name = "tika" }, | ||||
| { name = "tiktoken" }, | { name = "tiktoken" }, | ||||
| { name = "torch" }, | |||||
| { name = "umap-learn" }, | { name = "umap-learn" }, | ||||
| { name = "valkey" }, | { name = "valkey" }, | ||||
| { name = "vertexai" }, | { name = "vertexai" }, | ||||
| { name = "fastembed", marker = "platform_machine != 'x86_64' or sys_platform == 'darwin'" }, | { name = "fastembed", marker = "platform_machine != 'x86_64' or sys_platform == 'darwin'" }, | ||||
| { name = "fastembed-gpu", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, | { name = "fastembed-gpu", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, | ||||
| { name = "flagembedding" }, | { name = "flagembedding" }, | ||||
| { name = "torch" }, | |||||
| { name = "transformers" }, | { name = "transformers" }, | ||||
| ] | ] | ||||
| { name = "tencentcloud-sdk-python", specifier = "==3.0.1215" }, | { name = "tencentcloud-sdk-python", specifier = "==3.0.1215" }, | ||||
| { name = "tika", specifier = "==2.6.0" }, | { name = "tika", specifier = "==2.6.0" }, | ||||
| { name = "tiktoken", specifier = "==0.7.0" }, | { name = "tiktoken", specifier = "==0.7.0" }, | ||||
| { name = "torch", specifier = ">=2.5.0,<3.0.0" }, | |||||
| { name = "torch", marker = "extra == 'full'", specifier = ">=2.5.0,<3.0.0" }, | |||||
| { name = "transformers", marker = "extra == 'full'", specifier = ">=4.35.0,<5.0.0" }, | { name = "transformers", marker = "extra == 'full'", specifier = ">=4.35.0,<5.0.0" }, | ||||
| { name = "umap-learn", specifier = "==0.5.6" }, | { name = "umap-learn", specifier = "==0.5.6" }, | ||||
| { name = "valkey", specifier = "==6.0.2" }, | { name = "valkey", specifier = "==6.0.2" }, |