### What problem does this PR solve? Added cuda_is_available ### Type of change - [x] Refactoringtags/v0.16.0
| @@ -27,7 +27,6 @@ from . import operators | |||
| import math | |||
| import numpy as np | |||
| import cv2 | |||
| import torch | |||
| import onnxruntime as ort | |||
| from .postprocess import build_post_process | |||
| @@ -72,6 +71,15 @@ def load_model(model_dir, nm): | |||
| raise ValueError("not find model file path {}".format( | |||
| 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.enable_cpu_mem_arena = False | |||
| options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL | |||
| @@ -81,7 +89,7 @@ def load_model(model_dir, nm): | |||
| # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580 | |||
| # Shrink GPU memory after execution | |||
| run_options = ort.RunOptions() | |||
| if torch.cuda.is_available(): | |||
| if cuda_is_available(): | |||
| cuda_provider_options = { | |||
| "device_id": 0, # Use specific GPU | |||
| "gpu_mem_limit": 512 * 1024 * 1024, # Limit gpu memory | |||
| @@ -21,7 +21,6 @@ import numpy as np | |||
| import cv2 | |||
| from copy import deepcopy | |||
| import torch | |||
| import onnxruntime as ort | |||
| from huggingface_hub import snapshot_download | |||
| @@ -60,11 +59,21 @@ class Recognizer(object): | |||
| if not os.path.exists(model_file_path): | |||
| raise ValueError("not find model file path {}".format( | |||
| 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 | |||
| # Shrink GPU memory after execution | |||
| self.run_options = ort.RunOptions() | |||
| if torch.cuda.is_available(): | |||
| if cuda_is_available(): | |||
| options = ort.SessionOptions() | |||
| options.enable_cpu_mem_arena = False | |||
| cuda_provider_options = { | |||
| @@ -100,7 +100,6 @@ dependencies = [ | |||
| "tencentcloud-sdk-python==3.0.1215", | |||
| "tika==2.6.0", | |||
| "tiktoken==0.7.0", | |||
| "torch>=2.5.0,<3.0.0", | |||
| "umap_learn==0.5.6", | |||
| "vertexai==1.64.0", | |||
| "volcengine==1.0.146", | |||
| @@ -132,5 +131,6 @@ full = [ | |||
| "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'", | |||
| "flagembedding==1.2.10", | |||
| "torch>=2.5.0,<3.0.0", | |||
| "transformers>=4.35.0,<5.0.0" | |||
| ] | |||
| @@ -4814,7 +4814,6 @@ dependencies = [ | |||
| { name = "tencentcloud-sdk-python" }, | |||
| { name = "tika" }, | |||
| { name = "tiktoken" }, | |||
| { name = "torch" }, | |||
| { name = "umap-learn" }, | |||
| { name = "valkey" }, | |||
| { name = "vertexai" }, | |||
| @@ -4837,6 +4836,7 @@ full = [ | |||
| { name = "fastembed", marker = "platform_machine != 'x86_64' or sys_platform == 'darwin'" }, | |||
| { name = "fastembed-gpu", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, | |||
| { name = "flagembedding" }, | |||
| { name = "torch" }, | |||
| { name = "transformers" }, | |||
| ] | |||
| @@ -4946,7 +4946,7 @@ requires-dist = [ | |||
| { name = "tencentcloud-sdk-python", specifier = "==3.0.1215" }, | |||
| { name = "tika", specifier = "==2.6.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 = "umap-learn", specifier = "==0.5.6" }, | |||
| { name = "valkey", specifier = "==6.0.2" }, | |||