Browse Source

fix: 修复设备选择逻辑,确保 GPU 环境下使用 gpu:0

- 修改 _get_paddle_ocr_device() 函数
- 未配置设备时自动检测 CUDA 可用性
- NVIDIA GPU 环境默认返回 gpu:0
- NPU 环境返回 npu:0
- 避免错误使用 cpu 导致性能下降
何文松 1 ngày trước cách đây
mục cha
commit
7c129674c0
1 tập tin đã thay đổi với 11 bổ sung0 xóa
  1. 11 0
      pdf_converter_v2/utils/paddleocr_fallback.py

+ 11 - 0
pdf_converter_v2/utils/paddleocr_fallback.py

@@ -164,6 +164,17 @@ def _get_paddle_ocr_device() -> str:
     """获取 PaddleOCR Python API 使用的设备字符串(如 'gpu:0' 或 'cpu')"""
     devices = _get_paddle_ocr_devices()
     if not devices:
+        # 如果没有配置设备,根据环境自动选择
+        from .device_env import is_npu
+        if is_npu():
+            return "npu:0"
+        # NVIDIA GPU 环境,默认使用 gpu:0
+        try:
+            import torch
+            if torch.cuda.is_available():
+                return "gpu:0"
+        except ImportError:
+            pass
         return "cpu"
     global _PADDLE_OCR_DEVICE_INDEX
     with _PADDLE_OCR_DEVICE_LOCK: