diff --git a/api.py b/api.py index 87bda21..a615385 100644 --- a/api.py +++ b/api.py @@ -1,7 +1,19 @@ import logging +from config.env_setup import setup_environment + +# 환경 변수 설정을 최우선으로 호출 +setup_environment() + +# 로깅 기본 설정 +logging.basicConfig( + level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s - %(message)s" +) +logger = logging.getLogger("startup") + from fastapi import FastAPI from router import deepseek_router +from services.ocr_engine import init_engine logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s - %(message)s" @@ -14,6 +26,17 @@ app = FastAPI( ) +@app.on_event("startup") +async def startup_event(): + """FastAPI startup event handler.""" + logging.info("Application startup...") + try: + await init_engine() + logging.info("vLLM engine initialized successfully.") + except Exception as e: + logging.error(f"vLLM engine init failed: {e}", exc_info=True) + + @app.get("/health/API", include_in_schema=False) async def health_check(): return {"status": "API ok"} diff --git a/services/ocr_engine.py b/services/ocr_engine.py index 70107cd..bf5c7c9 100644 --- a/services/ocr_engine.py +++ b/services/ocr_engine.py @@ -1,37 +1,45 @@ import asyncio import io -from typing import Union +import logging import fitz from config.model_settings import CROP_MODE, MODEL_PATH, PROMPT -from fastapi import UploadFile from PIL import Image -from process.image_process import DeepseekOCRProcessor from vllm import AsyncLLMEngine, SamplingParams from vllm.engine.arg_utils import AsyncEngineArgs from vllm.model_executor.models.registry import ModelRegistry from services.deepseek_ocr import DeepseekOCRForCausalLM +from services.process.image_process import DeepseekOCRProcessor + +logger = logging.getLogger(__name__) # -------------------------------------------------------------------------- # 1. 모델 및 프로세서 초기화 # -------------------------------------------------------------------------- -# VLLM이 커스텀 모델을 인식하도록 등록 ModelRegistry.register_model("DeepseekOCRForCausalLM", DeepseekOCRForCausalLM) +_engine = None + + +async def init_engine(): + """vLLM 엔진을 비동기적으로 초기화합니다.""" + global _engine + if _engine is not None: + return + + engine_args = AsyncEngineArgs( + model=MODEL_PATH, + hf_overrides={"architectures": ["DeepseekOCRForCausalLM"]}, + block_size=256, + max_model_len=8192, + enforce_eager=False, + trust_remote_code=True, + tensor_parallel_size=1, + gpu_memory_utilization=0.75, + ) + _engine = AsyncLLMEngine.from_engine_args(engine_args) -# VLLM 비동기 엔진 설정 -engine_args = AsyncEngineArgs( - model=MODEL_PATH, - hf_overrides={"architectures": ["DeepseekOCRForCausalLM"]}, - block_size=256, - max_model_len=8192, - enforce_eager=False, - trust_remote_code=True, - tensor_parallel_size=1, - gpu_memory_utilization=0.75, -) -engine = AsyncLLMEngine.from_engine_args(engine_args) # 샘플링 파라미터 설정 sampling_params = SamplingParams( @@ -47,8 +55,11 @@ processor = DeepseekOCRProcessor() # 2. 핵심 처리 함수 # -------------------------------------------------------------------------- + async def _process_single_image(image: Image.Image) -> str: """단일 PIL 이미지를 받아 OCR을 수행하고 텍스트를 반환합니다.""" + if _engine is None: + raise RuntimeError("vLLM engine not initialized yet") if "" not in PROMPT: raise ValueError("프롬프트에 '' 토큰이 없어 OCR을 수행할 수 없습니다.") @@ -60,12 +71,13 @@ async def _process_single_image(image: Image.Image) -> str: request_id = f"request-{asyncio.get_running_loop().time()}" final_output = "" - async for request_output in engine.generate(request, sampling_params, request_id): + async for request_output in _engine.generate(request, sampling_params, request_id): if request_output.outputs: final_output = request_output.outputs[0].text return final_output + def _pdf_to_images(pdf_bytes: bytes, dpi=144) -> list[Image.Image]: """PDF 바이트를 받아 페이지별 PIL 이미지 리스트를 반환합니다.""" images = [] @@ -83,6 +95,7 @@ def _pdf_to_images(pdf_bytes: bytes, dpi=144) -> list[Image.Image]: pdf_document.close() return images + async def process_document(file_bytes: bytes, content_type: str, filename: str) -> dict: """ 업로드된 파일(이미지 또는 PDF)을 처리하여 OCR 결과를 반환합니다. @@ -114,4 +127,4 @@ async def process_document(file_bytes: bytes, content_type: str, filename: str) raise ValueError( f"지원하지 않는 파일 형식입니다: {content_type}. " "이미지(JPEG, PNG 등) 또는 PDF 파일을 업로드해주세요." - ) \ No newline at end of file + )