ocr_platform/ocr_utils/device_utils.py

"""
设备检测工具
从 MinerU 移植，用于自动检测可用设备 (CUDA/MPS/NPU/CPU)
"""
import os

try:
    import torch
except ImportError:
    torch = None

try:
    import torch_npu
except ImportError:
    torch_npu = None


def get_device():
    """
    自动检测并返回可用的设备
    
    优先级: CUDA > MPS > NPU > CPU
    
    Returns:
        str: 设备名称 ('cuda', 'mps', 'npu', 'cpu')
    
    Environment Variables:
        MINERU_DEVICE_MODE: 强制指定设备模式
    """
    # 支持通过环境变量强制指定设备
    device_mode = os.getenv('MINERU_DEVICE_MODE', None)
    if device_mode is not None:
        return device_mode
    
    # 如果没有 torch，返回 cpu
    if torch is None:
        return "cpu"
    
    # 自动检测
    if torch.cuda.is_available():
        return "cuda"
    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        return "mps"
    else:
        # 尝试检测华为 NPU
        try:
            if torch_npu is not None and torch_npu.npu.is_available():
                return "npu"
        except Exception:
            pass
    
    return "cpu"


def get_device_name():
    """
    获取设备的友好名称
    
    Returns:
        str: 设备的友好名称
    """
    device = get_device()
    
    device_names = {
        "cuda": "NVIDIA CUDA",
        "mps": "Apple Metal (MPS)",
        "npu": "Huawei NPU",
        "cpu": "CPU"
    }
    
    return device_names.get(device, device.upper())


if __name__ == "__main__":
    """测试设备检测"""
    print(f"🔍 Detecting available device...")
    device = get_device()
    device_name = get_device_name()
    print(f"✅ Device: {device} ({device_name})")
    
    # 测试 torch 设备
    if torch is not None:
        try:
            test_tensor = torch.tensor([1.0, 2.0, 3.0])
            if device != "cpu":
                test_tensor = test_tensor.to(device)
                print(f"✅ Torch tensor moved to {device}")
            print(f"   Tensor device: {test_tensor.device}")
        except Exception as e:
            print(f"⚠️  Failed to move tensor to {device}: {e}")