#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ DeepSeek OCR API Server (vLLM) - 极简版 + 优化版 """ import os import io import re import argparse import asyncio from io import BytesIO from typing import List from concurrent.futures import ThreadPoolExecutor import torch from PIL import Image try: import fitz except Exception: fitz = None from fastapi import FastAPI, File, UploadFile, HTTPException, Form from fastapi.responses import JSONResponse from fastapi.middleware.cors import CORSMiddleware import uvicorn from vllm import LLM, SamplingParams from vllm.model_executor.models.registry import ModelRegistry from deepseek_ocr import DeepseekOCRForCausalLM from process.ngram_norepeat import NoRepeatNGramLogitsProcessor from process.image_process import DeepseekOCRProcessor app = FastAPI(title="DeepSeek OCR API (vLLM) - Optimized", version="2.0.0") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) llm = None cpu_executor = ThreadPoolExecutor(max_workers=8, thread_name_prefix="CPU-Worker") gpu_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="GPU-Worker") vllm_lock = asyncio.Lock() PROMPT_OCR = "\n<|grounding|>Convert the document to markdown." PROMPT_DESC = "\nDescribe this image in detail." # ----------------------- # Monkey Patch # ----------------------- _original_tokenize = DeepseekOCRProcessor.tokenize_with_images def _patched_tokenize(self, images, bos=True, eos=True, cropping=True, prompt=None): if prompt is not None: import config old = config.PROMPT config.PROMPT = prompt try: return _original_tokenize(self, images, bos, eos, cropping) finally: config.PROMPT = old return _original_tokenize(self, images, bos, eos, cropping) DeepseekOCRProcessor.tokenize_with_images = _patched_tokenize def pdf_to_images_sync(pdf_bytes: bytes, dpi: int = 144) -> List[Image.Image]: """PDF 转图片 """ if fitz is None: raise RuntimeError("Please install PyMuPDF") images = [] doc = fitz.open(stream=pdf_bytes, filetype="pdf") matrix = fitz.Matrix(dpi / 72.0, dpi / 72.0) for page in doc: pix = page.get_pixmap(matrix=matrix, alpha=False) img = Image.open(io.BytesIO(pix.tobytes("png"))) if img.mode != "RGB": if img.mode in ("RGBA", "LA"): bg = Image.new("RGB", img.size, (255, 255, 255)) bg.paste(img, mask=img.split()[-1]) img = bg else: img = img.convert("RGB") images.append(img) doc.close() return images def image_open_sync(image_bytes: bytes) -> Image.Image: """打开图片 (同步版本)""" return Image.open(BytesIO(image_bytes)).convert("RGB") def clear_vllm_cache_sync(): """清理 vLLM 缓存 (同步版本)""" if llm is None: return try: if hasattr(llm.llm_engine, 'input_preprocessor'): prep = llm.llm_engine.input_preprocessor if hasattr(prep, '_mm_processor_cache'): prep._mm_processor_cache.clear() except: pass def tokenize_image_sync(image: Image.Image, prompt: str): """ 图像 tokenize (同步版本, CPU 密集) WARNING: 这是最大的优化点! """ processor = DeepseekOCRProcessor() return processor.tokenize_with_images(images=[image], prompt=prompt) def vllm_generate_sync(tokenized, prompt: str) -> str: """ vLLM 推理 (同步版本, GPU 密集) 注意: tokenized 已经在 CPU 线程池完成 """ batch_inputs = [{ "prompt": prompt, "multi_modal_data": {"image": tokenized} }] if prompt == PROMPT_OCR: logits_proc = [NoRepeatNGramLogitsProcessor(20, 50, {128821, 128822})] params = SamplingParams( temperature=0.0, max_tokens=8192, skip_special_tokens=False, logits_processors=logits_proc, repetition_penalty=1.05, ) else: params = SamplingParams( temperature=0.0, max_tokens=8192, skip_special_tokens=False, ) outputs = llm.generate(batch_inputs, params) return outputs[0].outputs[0].text def clean_markdown_sync(text: str) -> str: """清理 Markdown (同步版本)""" text = re.sub(r'<\|ref\|>.*?<\|/ref\|>', '', text) text = re.sub(r'<\|det\|>.*?<\|/det\|>', '', text) text = re.sub(r'<\|.*?\|>', '', text) text = re.sub(r'\[\[.*?\]\]', '', text) text = re.sub(r'={50,}.*?={50,}', '', text, flags=re.DOTALL) text = re.sub(r'\n{3,}', '\n\n', text) return text.strip() async def pdf_to_images_async(pdf_bytes: bytes, dpi: int = 144) -> List[Image.Image]: """PDF 转图片 (异步)""" loop = asyncio.get_event_loop() return await loop.run_in_executor(cpu_executor, pdf_to_images_sync, pdf_bytes, dpi) async def image_open_async(image_bytes: bytes) -> Image.Image: """打开图片 (异步)""" loop = asyncio.get_event_loop() return await loop.run_in_executor(cpu_executor, image_open_sync, image_bytes) async def tokenize_image_async(image: Image.Image, prompt: str): """ 图像 tokenize (异步) NOTE: 关键优化: 在 CPU 线程池执行 """ loop = asyncio.get_event_loop() return await loop.run_in_executor(cpu_executor, tokenize_image_sync, image, prompt) async def vllm_generate_async(image: Image.Image, prompt: str) -> str: """ 完整的 vLLM 推理流程 (异步) 优化: 分离 tokenize (CPU) 和 generate (GPU) """ # 步骤1: tokenize (CPU 密集, 在 CPU 线程池执行) tokenized = await tokenize_image_async(image, prompt) # 步骤2: GPU 推理 (GPU 密集, 在 GPU 线程池执行, 有锁保护) async with vllm_lock: # 清理缓存 (在 GPU 线程池执行) loop = asyncio.get_event_loop() await loop.run_in_executor(gpu_executor, clear_vllm_cache_sync) # GPU 推理 result = await loop.run_in_executor( gpu_executor, vllm_generate_sync, tokenized, prompt ) return result async def clean_markdown_async(text: str) -> str: """清理 Markdown (异步)""" loop = asyncio.get_event_loop() return await loop.run_in_executor(cpu_executor, clean_markdown_sync, text) async def generate_image_description_async(image: Image.Image) -> str: """生成图片描述 (异步)""" try: # GPU 推理 result = await vllm_generate_async(image, PROMPT_DESC) # CPU 后处理 loop = asyncio.get_event_loop() def process_desc(text): desc = re.sub(r'<\|ref\|>.*?<\|/ref\|>', '', text) desc = re.sub(r'<\|det\|>.*?<\|/det\|>', '', desc) desc = re.sub(r'<\|.*?\|>', '', desc) desc = re.sub(r'\[\[.*?\]\]', '', desc) desc = re.sub(r'\s+', ' ', desc).strip() if len(desc) > 200: cutoff = desc[:200].rfind('.') if cutoff > 100: desc = desc[:cutoff + 1] else: desc = desc[:200].rsplit(' ', 1)[0] + '...' return desc desc = await loop.run_in_executor(cpu_executor, process_desc, result) return desc except Exception as e: print(f"WARNING: 图片描述失败: {e}") return "" # ----------------------- # 模型初始化 # ----------------------- def initialize_model(model_path: str, gpu_id: int): global llm ModelRegistry.register_model("DeepseekOCRForCausalLM", DeepseekOCRForCausalLM) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) os.environ['VLLM_USE_V1'] = '0' print(f"[INFO] 加载模型: {model_path}") llm = LLM( model=model_path, hf_overrides={"architectures": ["DeepseekOCRForCausalLM"]}, block_size=64, enforce_eager=False, trust_remote_code=True, max_model_len=8192, tensor_parallel_size=1, gpu_memory_utilization=0.9, max_num_seqs=20, disable_mm_preprocessor_cache=True, ) print("[SUCCESS] 模型加载完成") #print(f"[INFO] 线程池配置:") #print(f" - CPU 线程池: {cpu_executor._max_workers} 线程") #print(f" - GPU 线程池: {gpu_executor._max_workers} 线程") # ----------------------- # API 路由 # ----------------------- @app.get("/") async def root(): return { "service": "DeepSeek OCR (vLLM) - Optimized", "version": "2.0.0", "status": "running" } @app.get("/health") async def health(): return { "status": "healthy", "model_ready": llm is not None, "cpu_workers": cpu_executor._max_workers, "gpu_workers": gpu_executor._max_workers, } async def vllm_generate_batch_async( images: List[Image.Image], prompt: str, show_progress: bool = True ) -> List[str]: """ 批量 vLLM 推理 - 真正的批处理优化 Args: images: 图片列表 prompt: 提示词 show_progress: 是否显示进度 Returns: 生成的文本列表 """ total = len(images) # 步骤1: 并发 tokenize # 标准化图片 -> Vision Encoder (ViT) -> 图像特征向量 (例如：[196, 1024] - 196个位置，每个1024维) if show_progress: print(f" [1/3] Tokenize {total} 页...") tokenize_tasks = [tokenize_image_async(img, prompt) for img in images] all_tokenized = await asyncio.gather(*tokenize_tasks) if show_progress: print(f" [1/3] Tokenize 完成") # 步骤2: 构造批量输入 batch_inputs = [ { "prompt": prompt, "multi_modal_data": {"image": tok} } for tok in all_tokenized ] # 步骤3: 批量 GPU 推理 async with vllm_lock: if show_progress: print(f" [2/3] GPU 批量推理 {total} 页...") loop = asyncio.get_event_loop() # 清理缓存 await loop.run_in_executor(gpu_executor, clear_vllm_cache_sync) # 批量推理 def batch_generate(): # 根据 prompt 类型选择参数 if prompt == PROMPT_OCR: logits_proc = [NoRepeatNGramLogitsProcessor(20, 50, {128821, 128822})] params = SamplingParams( temperature=0.0, max_tokens=8192, skip_special_tokens=False, logits_processors=logits_proc, repetition_penalty=1.05, ) else: params = SamplingParams( temperature=0.0, max_tokens=8192, skip_special_tokens=False, ) # NOTE: 关键: 批量调用 outputs = llm.generate(batch_inputs, params) return [out.outputs[0].text for out in outputs] results = await loop.run_in_executor(gpu_executor, batch_generate) if show_progress: print(f" [2/3] GPU 推理完成") return results @app.post("/ocr") async def ocr( file: UploadFile = File(...), enable_description: bool = Form(False), ): """OCmR 接口 (批量处理)""" if llm is None: raise HTTPException(503, "模型未加载") import time start_time = time.time() try: # 1. 读取文件 contents = await file.read() t1 = time.time() # 2. 解析文件 if file.filename.lower().endswith('.pdf'): # 如果是PDF，则转换为图片列表 images = await pdf_to_images_async(contents) else: # 如果是图片，则直接打开 images = [await image_open_async(contents)] t2 = time.time() # 3. 批量 OCR raw_results = await vllm_generate_batch_async(images, PROMPT_OCR) t3 = time.time() print(f" OCR 耗时: {t3 - t2:.2f}s") # 4. 后处理 print(f" [3/3] 后处理...") async def postprocess(idx: int, raw: str, img: Image.Image) -> str: # 图片描述 if enable_description: img_pattern = r'<\|ref\|>image<\|/ref\|><\|det\|>\[\[.*?\]\]<\|/det\|>' matches = list(re.finditer(img_pattern, raw)) for match in matches: desc = await generate_image_description_async(img) replacement = f"[图片: {desc}]" if desc else "[图片]" raw = raw.replace(match.group(0), replacement) # 清理 Markdown cleaned = await clean_markdown_async(raw) return cleaned if cleaned else "" tasks = [postprocess(i, raw, img) for i, (raw, img) in enumerate(zip(raw_results, images))] md_parts = await asyncio.gather(*tasks) t4 = time.time() print(f" [3/3] 后处理完成 ({t4 - t3:.2f}s)") # 5. 合并结果 final_md = "\n\n".join([md for md in md_parts if md]) total_time = time.time() - start_time print(f"{'='*60}") print(f"[SUCCESS] 全部完成") print(f" 总耗时: {total_time:.2f}s") print(f" 平均: {total_time / len(images):.2f}s/页") print(f"{'='*60}\n") return JSONResponse({ "markdown": final_md, "page_count": len(images), "processing_time": round(total_time, 2), }) except Exception as e: import traceback print(f"[ERROR] 处理失败: {e}") print(traceback.format_exc()) raise HTTPException(500, f"处理失败: {e}") # ----------------------- # 优雅关闭 # ----------------------- @app.on_event("shutdown") async def shutdown_event(): print("[INFO] 关闭线程池...") cpu_executor.shutdown(wait=True) gpu_executor.shutdown(wait=True) print("[SUCCESS] 线程池已关闭") # ----------------------- # 启动 # ----------------------- def main(): parser = argparse.ArgumentParser() parser.add_argument("--model-path", required=True, help="模型路径") parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID") parser.add_argument("--port", type=int, default=8002, help="端口") parser.add_argument("--host", default="0.0.0.0", help="监听地址") parser.add_argument("--cpu-workers", type=int, default=2, help="CPU 线程池大小") args = parser.parse_args() # 更新线程池大小 global cpu_executor cpu_executor = ThreadPoolExecutor( max_workers=args.cpu_workers, thread_name_prefix="CPU-Worker" ) initialize_model(args.model_path, args.gpu_id) print(f"\n[INFO] 服务启动: http://{args.host}:{args.port}") print(f"[INFO] 接口文档: http://{args.host}:{args.port}/docs\n") uvicorn.run(app, host=args.host, port=args.port, workers=1) if __name__ == "__main__": main()