From e1046ba3faeb0877400e7cb2b2bd24e3ad020ff3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=89=E6=B4=8B=E4=B8=89=E6=B4=8B?= <1258009915@qq.com> Date: Wed, 7 Aug 2024 01:19:26 +0800 Subject: [PATCH] Change to better import dependency --- src/models/ocr_model/model/TexTeller.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/ocr_model/model/TexTeller.py b/src/models/ocr_model/model/TexTeller.py index 0a811f2..1f7e0ac 100644 --- a/src/models/ocr_model/model/TexTeller.py +++ b/src/models/ocr_model/model/TexTeller.py @@ -1,5 +1,4 @@ from pathlib import Path -from optimum.onnxruntime import ORTModelForVision2Seq from ...globals import ( VOCAB_SIZE, @@ -32,6 +31,7 @@ class TexTeller(VisionEncoderDecoderModel): if not use_onnx: return VisionEncoderDecoderModel.from_pretrained(cls.REPO_NAME) else: + from optimum.onnxruntime import ORTModelForVision2Seq use_gpu = True if onnx_provider == 'cuda' else False return ORTModelForVision2Seq.from_pretrained(cls.REPO_NAME, provider="CUDAExecutionProvider" if use_gpu else "CPUExecutionProvider") model_path = Path(model_path).resolve()