Commit f5112e21 authored by myhloli's avatar myhloli
Browse files

fix(model): move environment variable settings to global scope

- Move environment variable settings for NPU, MPS, and other configurations to the global scope in doc_analyze_by_custom_model.py
- Remove redundant environment variable settings in pdf_extract_kit.py
- This change ensures consistent configuration across the application and avoids potential conflicts or duplicate settings
parent 4021abeb
import os import os
import time import time
import torch import torch
os.environ['FLAGS_npu_jit_compile'] = '0' # 关闭paddle的jit编译
os.environ['FLAGS_use_stride_kernel'] = '0'
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1' # 让mps可以fallback
os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新
# 关闭paddle的信号处理 # 关闭paddle的信号处理
import paddle import paddle
paddle.disable_signal_handler() paddle.disable_signal_handler()
...@@ -12,11 +15,8 @@ from loguru import logger ...@@ -12,11 +15,8 @@ from loguru import logger
from magic_pdf.model.batch_analyze import BatchAnalyze from magic_pdf.model.batch_analyze import BatchAnalyze
from magic_pdf.model.sub_modules.model_utils import get_vram from magic_pdf.model.sub_modules.model_utils import get_vram
os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新
try: try:
import torchtext import torchtext
if torchtext.__version__ >= '0.18.0': if torchtext.__version__ >= '0.18.0':
torchtext.disable_torchtext_deprecation_warning() torchtext.disable_torchtext_deprecation_warning()
except ImportError: except ImportError:
......
...@@ -89,13 +89,6 @@ class CustomPEKModel: ...@@ -89,13 +89,6 @@ class CustomPEKModel:
# 初始化解析方案 # 初始化解析方案
self.device = kwargs.get('device', 'cpu') self.device = kwargs.get('device', 'cpu')
if str(self.device).startswith("npu"):
import torch_npu
os.environ['FLAGS_npu_jit_compile'] = '0'
os.environ['FLAGS_use_stride_kernel'] = '0'
elif str(self.device).startswith("mps"):
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
logger.info('using device: {}'.format(self.device)) logger.info('using device: {}'.format(self.device))
models_dir = kwargs.get( models_dir = kwargs.get(
'models_dir', os.path.join(root_dir, 'resources', 'models') 'models_dir', os.path.join(root_dir, 'resources', 'models')
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment