GOT-OCR-2.0-hf 昇腾 NPU 部署实战课程

┌──────────────────────────────────────────────────────┐
│              FastAPI + Uvicorn 服务层                  │
│  ┌────────────────────────────────────────────────┐    │
│  │  main:app (port 8016)                          │    │
│  │  ├─ /           → 健康检查                      │    │
│  │  └─ /v1/orc/stepFun → OCR 推理接口（文件上传） │    │
│  └──────────────────┬───────────────────────────────┘    │
│                     │                                 │
│  ┌──────────────────▼───────────────────────────┐    │
│  │  stepFunService.py（核心推理链路）             │    │
│  │  ├─ stepFunModel() → 单例模型加载             │    │
│  │  ├─ processor(image) → 图像编码 + 文本编码    │    │
│  │  ├─ model.generate() → 自回归 OCR 解码       │    │
│  │  └─ regex 提取 → OCR 文本清洗                 │    │
│  └──────────────────┬───────────────────────────┘    │
│                     │                                 │
├─────────────────────┼────────────────────────────────┤
│       torch_npu + CANN 8.5                           │
│  ┌──────────────┐  ┌─▼──────────────┐                │
│  │  GE 图编译   │  │ CUBE + Vector  │                │
│  │  (算子融合)  │  │  混合执行      │                │
│  └──────────────┘  └────────────────┘                │
├──────────────────────────────────────────────────────┤
│              昇腾 NPU 硬件层                          │
│  ┌──────────┐ ┌──────────┐ ┌──────────────┐          │
│  │ Scalar   │ │ Vector   │ │ CUBE 矩阵乘  │          │
│  │ 控制流   │ │ 图像编码 │ │ 文本解码     │          │
│  └──────────┘ └──────────┘ └──────────────┘          │
└──────────────────────────────────────────────────────┘

# [依赖安装] 一条命令安装所有 Python 依赖
!pip install -U atomgit torch_npu transformers accelerate fastapi uvicorn python-multipart pillow -i https://mirrors.huaweicloud.com/repository/pypi/simple

# [代码获取] 从 AtomGit 克隆 GOT-OCR-2.0-hf 服务代码
!mkdir -p /opt/atomgit
!git clone https://atomgit.com/atomgit-ascend/GOT-OCR-2.0-hf.git /opt/atomgit/GOT-OCR-2.0-hf

# [模型下载] 从 AtomGit 平台下载 GOT-OCR-2.0-hf 模型
from atomgit_hub import snapshot_download

model_repo = "StepFun/GOT-OCR-2.0-hf"
local_dir = "/opt/atomgit/models/GOT-OCR-2.0-hf"

snapshot_download(model_repo, local_dir=local_dir)
print(f"模型已下载到: {local_dir}")

# [核心适配] stepFunService.py 完整替换内容
# 写入路径：/opt/atomgit/GOT-OCR-2.0-hf/app/service/stepFunService.py

# ===== 强制 NPU 设置 - 必须在最开头 =====
import os
os.environ['NPU_VISIBLE_DEVICES'] = '0'
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
os.environ['ASCEND_DEVICE_ID'] = '0'
os.environ['DEVICE_ID'] = '0'
os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'
# ========================================

import sys
import traceback
import time
import re

from utils.loggingConfig import get_logger

logger = get_logger('')

# ========== NPU 设备配置 ==========
import torch
import torch_npu

device = 'cpu'
devicemap = 'cpu'

try:
    if torch_npu.npu.is_available():
        device = 'npu:0'
        devicemap = device
        logger.info(f'✅ 检测到 NPU，使用设备: {device}')
        torch.npu.set_device(0)
    else:
        logger.warning('⚠️ NPU 不可用，使用 CPU')
except Exception as e:
    logger.error(f'NPU 检测异常: {e}')
    device = 'cpu'
    devicemap = 'cpu'

logger.info(f'最终设备配置 - device: {device}, devicemap: {devicemap}')
# ===============================

from transformers import AutoProcessor, AutoModelForImageTextToText

class stepFunModel():
    _instance = None
    _initialized = False

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(stepFunModel, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if stepFunModel._initialized:
            return

        model_path = os.getenv('STEPFUN_MODEL_PATH')
        logger.info(f'模型路径: {model_path}')
        logger.info(f'使用的 device_map: {devicemap}')

        try:
            model = AutoModelForImageTextToText.from_pretrained(
                model_path,
                device_map=devicemap,
                torch_dtype=torch.bfloat16 if device != 'cpu' else torch.float32
            )
            processor = AutoProcessor.from_pretrained(model_path)

            self._model = model
            self._processor = processor
            stepFunModel._initialized = True

            actual_device = next(model.parameters()).device
            logger.info(f'✅ 模型实际加载设备: {actual_device}')
            logger.info(f'模型加载完成')
        except Exception as e:
            logger.error(f'模型加载失败: {e}')
            traceback.print_exc()
            raise

stepFunModel()

async def stepFunService(image_path):
    logger.info(f"处理图片: {image_path}")
    t1 = time.time()
    model = stepFunModel()

    inputs = model._processor(image_path, return_tensors="pt").to(device)

    generate_ids = model._model.generate(
        **inputs,
        do_sample=False,
        tokenizer=model._processor.tokenizer,
        stop_strings="<|im_end|>",
        max_new_tokens=4096,
    )

    full_output = model._processor.decode(generate_ids[0], skip_special_tokens=False)

    # 从 assistant 标签中提取 OCR 文本
    pattern = r'<\|im_start\|>assistant\n(.*?)<\|im_end\|>'
    match = re.search(pattern, full_output, re.DOTALL)

    if match:
        ocr_result = match.group(1).strip()
    else:
        ocr_result = re.sub(r'<\|.*?\|>', '', full_output).strip()

    t2 = time.time()
    logger.info(f'OCR 耗时: {t2 - t1:.2f}秒')

    return ocr_result

# [控制器适配] modelExperienceController.py 完整替换内容
# 写入路径：/opt/atomgit/GOT-OCR-2.0-hf/app/controller/modelExperienceController.py

import os
import time
from fastapi import APIRouter, UploadFile, File, Form
from typing import Optional
import shutil
from service.stepFunService import stepFunService
from utils.loggingConfig import get_logger

logger = get_logger('')
router = APIRouter()

ALLOWED_EXTENSIONS = os.getenv('ALLOWED_EXTENSIONS', '.tif,.png,.webp,.gif,.jpeg,.jpg,.bmp,.tiff').split(',')

@router.post("/v1/orc/stepFun")
async def stepFun(
    file: UploadFile = File(...),
    model: Optional[str] = Form(None)
):
    logger.info(f"######### stepFun")
    logger.info(f"########### file {file.filename}")

    file_ext = os.path.splitext(file.filename)[1].lower()
    if file_ext not in ALLOWED_EXTENSIONS:
        return {
            "success": False,
            "code": 400,
            "error_msg": f"不支持的文件格式: {file_ext}"
        }

    temp_dir = "/tmp/ocr_uploads"
    os.makedirs(temp_dir, exist_ok=True)
    temp_path = os.path.join(temp_dir, f"{int(time.time())}_{file.filename}")

    try:
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        ocr_result = await stepFunService(temp_path)

        return {
            "success": True,
            "code": 200,
            "ocr_msg": ocr_result if ocr_result else ""
        }

    except Exception as e:
        logger.error(f"OCR 处理失败: {e}")
        return {
            "success": False,
            "code": 500,
            "error_msg": str(e)
        }
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)

#!/bin/bash

echo "=========================================="
echo "启动 GOT-OCR-2.0-hf 服务 (uvicorn)"
echo "=========================================="

# 清除冲突变量
unset ASCEND_VISIBLE_DEVICES
unset CUDA_VISIBLE_DEVICES

# 设置环境变量
export NPU_VISIBLE_DEVICES=0
export ASCEND_RT_VISIBLE_DEVICES=0
export ASCEND_DEVICE_ID=0
export DEVICE_ID=0
export STEPFUN_MODEL_PATH=/opt/atomgit/models/GOT-OCR-2.0-hf
export TORCH_DEVICE_BACKEND_AUTOLOAD=0

echo "环境变量:"
echo "  NPU_VISIBLE_DEVICES = $NPU_VISIBLE_DEVICES"
echo "  STEPFUN_MODEL_PATH = $STEPFUN_MODEL_PATH"

# 清理 Python 缓存
find /opt/atomgit/GOT-OCR-2.0-hf -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null
find /opt/atomgit/GOT-OCR-2.0-hf -name "*.pyc" -delete

echo ""
echo "启动 uvicorn 服务..."
echo "服务地址: http://0.0.0.0:8016"
echo "OCR 接口: curl -X POST http://localhost:8016/v1/orc/stepFun -F 'file=@image.jpg'"
echo "=========================================="

cd /opt/atomgit/GOT-OCR-2.0-hf/app
python -m uvicorn main:app --host 0.0.0.0 --port 8016 --log-level info

# [服务启动] 执行启动脚本
!chmod +x /opt/atomgit/GOT-OCR-2.0-hf/app/run.sh
!/opt/atomgit/GOT-OCR-2.0-hf/app/run.sh

# [OCR 推理] 使用仓库自带的测试图片验证
!curl -X POST http://localhost:8016/v1/orc/stepFun \
  -F "file=@/opt/atomgit/GOT-OCR-2.0-hf/t-unit/image_ocr.jpg" \
  -F "model=StepFun/GOT-OCR-2.0-hf"

# ✅ 正确：环境变量在 import 之前
import os
os.environ['NPU_VISIBLE_DEVICES'] = '0'
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0'
import torch_npu

unset ASCEND_VISIBLE_DEVICES
unset CUDA_VISIBLE_DEVICES
export NPU_VISIBLE_DEVICES=0
export ASCEND_RT_VISIBLE_DEVICES=0

GET /

POST /v1/orc/stepFun
Content-Type: multipart/form-data

{
  "success": true,
  "code": 200,
  "ocr_msg": "识别出的文字内容"
}