113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
import os
|
|
import base64
|
|
import io
|
|
import time
|
|
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
from PIL import Image
|
|
|
|
# vLLM 및 모델 관련 import
|
|
from vllm import AsyncLLMEngine, SamplingParams
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
from vllm.model_executor.models.registry import ModelRegistry
|
|
|
|
# DeepSeek-OCR 관련 로컬 import
|
|
from deepseek_ocr import DeepseekOCRForCausalLM
|
|
from process.image_process import DeepseekOCRProcessor
|
|
from process.ngram_norepeat import NoRepeatNGramLogitsProcessor
|
|
|
|
# --- Configuration ---
|
|
# Docker 환경에서는 환경 변수를 사용하거나, Dockerfile에서 모델을 다운로드하는 것이 좋습니다.
|
|
# 여기서는 config.py의 기본값을 사용하되, 환경 변수로 재정의할 수 있도록 합니다.
|
|
MODEL_PATH = os.environ.get("MODEL_PATH", "deepseek-ai/deepseek-vl-7b-base")
|
|
# 참고: 실제 `config.py`는 로컬 경로를 사용하므로, 허깅페이스 모델 ID로 대체합니다.
|
|
# 이 모델을 사용하려면 인터넷 연결이 필요하며, 처음 실행 시 다운로드됩니다.
|
|
|
|
# --- Model Initialization ---
|
|
|
|
# 1. 커스텀 모델 등록
|
|
ModelRegistry.register_model("DeepseekOCRForCausalLM", DeepseekOCRForCausalLM)
|
|
|
|
# 2. vLLM 엔진 설정
|
|
engine_args = AsyncEngineArgs(
|
|
model=MODEL_PATH,
|
|
hf_overrides={"architectures": ["DeepseekOCRForCausalLM"]},
|
|
max_model_len=8192,
|
|
enforce_eager=False,
|
|
trust_remote_code=True,
|
|
tensor_parallel_size=1, # 단일 GPU 사용
|
|
gpu_memory_utilization=0.90, # GPU 메모리 사용률
|
|
)
|
|
engine = AsyncLLMEngine.from_engine_args(engine_args)
|
|
|
|
# 3. Deepseek OCR 프로세서 초기화
|
|
processor = DeepseekOCRProcessor()
|
|
|
|
# 4. FastAPI 앱 초기화
|
|
app = FastAPI()
|
|
|
|
# --- Pydantic Models ---
|
|
class InferenceRequest(BaseModel):
|
|
# Base64로 인코딩된 이미지 문자열
|
|
base64_image: str
|
|
|
|
class InferenceResponse(BaseModel):
|
|
text: str
|
|
|
|
# --- API Endpoints ---
|
|
|
|
@app.get("/")
|
|
def health_check():
|
|
return {"status": "DeepSeek-OCR service is running"}
|
|
|
|
@app.post("/process", response_model=InferenceResponse)
|
|
async def process_image(request: InferenceRequest):
|
|
"""
|
|
Base64 인코딩된 이미지를 받아 OCR 추론을 수행합니다.
|
|
"""
|
|
try:
|
|
# 1. Base64 이미지 디코딩
|
|
image_data = base64.b64decode(request.base64_image)
|
|
image = Image.open(io.BytesIO(image_data)).convert('RGB')
|
|
|
|
# 2. 이미지 전처리
|
|
prompt = "<image>"
|
|
image_features = processor.tokenize_with_images(
|
|
images=[image],
|
|
bos=True,
|
|
eos=True,
|
|
cropping=False # CROP_MODE 기본값 사용
|
|
)
|
|
|
|
# 3. 샘플링 파라미터 설정 (기존 스크립트 참조)
|
|
logits_processors = [NoRepeatNGramLogitsProcessor(ngram_size=30, window_size=90, whitelist_token_ids={128821, 128822})]
|
|
sampling_params = SamplingParams(
|
|
temperature=0.0,
|
|
max_tokens=8192,
|
|
logits_processors=logits_processors,
|
|
skip_special_tokens=False,
|
|
)
|
|
|
|
# 4. vLLM으로 추론 실행
|
|
request_id = f"dpsk-request-{int(time.time())}"
|
|
vllm_request = {
|
|
"prompt": prompt,
|
|
"multi_modal_data": {"image": image_features}
|
|
}
|
|
|
|
final_output = None
|
|
async for request_output in engine.generate(vllm_request, sampling_params, request_id):
|
|
# 스트리밍 결과의 마지막 최종본을 사용
|
|
final_output = request_output
|
|
|
|
if final_output and final_output.outputs:
|
|
generated_text = final_output.outputs[0].text
|
|
return InferenceResponse(text=generated_text)
|
|
else:
|
|
raise Exception("Model generated no output.")
|
|
|
|
except Exception as e:
|
|
# 실제 운영 환경에서는 로깅을 추가하는 것이 좋습니다.
|
|
return {"error": f"An error occurred: {str(e)}"}, 500
|