[FastAPI] VLM 서빙하기

FastAPI를 이용하여 로컬에서 VLM(Vision Language Model)을 서빙해보았습니다.

VLM은 image와 prompt를 입력으로 받아 text를 return 합니다.

 

1. 로컬에서 VLM 서빙

a. 라이브러리 설치

# 가상환경 생성
conda create -n vlm_api python=3.11

# 가상환경 실행
conda activate vlm_api

# 라이브러리 설치 for VLM
pip install torch transformers pillow requests swarms huggingface_hub

# 허깅페이스 로그인
huggingface-cli login

# 라이브러리 설치 for FastAPI
pip install fastapi uvicorn

 

b. api 파일 생성

저는 VLM으로 PaliGemma를 사용했지만 다른 모델을 사용하고 싶으시다면 ### Model Initialize ###와 ### Model Inference ### 부분을 수정해서 사용하시면 됩니다.

vlm_api.py 파일을 만들어 아래 코드를 붙여넣습니다.

from fastapi import FastAPI, File, UploadFile, Form
from pydantic import BaseModel
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from PIL import Image
import requests
import torch


# FastAPI 앱 생성
app = FastAPI()

##### Model Initialize #####
model_id = "google/paligemma-3b-pt-224"
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).eval()
processor = AutoProcessor.from_pretrained(model_id)
############################

# 이미지와 텍스트를 함께 처리하는 엔드포인트 설정
@app.post("/predict/")
async def predict(file: UploadFile = File(...), text_data: str = Form(None)):
    # 이미지를 PIL 형식으로 로드
    image = Image.open(file.file).convert('RGB')
    print(text_data)
    if text_data is not None: 
        # 입력받은 텍스트를 프롬프트로 사용
        prompt = text_data
    else: # 미리 설정한 프롬프트 사용
        prompt = """
        ocr name and price of the product.
        """

    model_inputs = processor(text=prompt, images=image, return_tensors="pt")
    input_len = model_inputs["input_ids"].shape[-1]

    ##### model Inference #####
    with torch.inference_mode():
        generation = model.generate(**model_inputs, max_new_tokens=100, do_sample=False)
        generation = generation[0][input_len:]
        decoded = processor.decode(generation, skip_special_tokens=True)
        print(decoded)
    ###########################
    
    # 결과 반환
    return {"result": decoded}

c. uvicorn 실행

uvicorn vlm_api:app --reload

 

d. test

127.0.0.1:8000/docs 링크로 들어가시면 Swagger를 통해 편하게 테스트 해볼 수 있습니다.

 

저는 바지를 이미지로 넣고, prompt를 'ocr name and price'로 설정했습니다.

결과는

MILITARY CARGO PANT HIPSTOP
₩245,000

 

RIPSTOP인데 HIPSTOP이라고 해서 살짝 틀렸긴하지만 모델 서빙은 제대로 되었네요ㅎㅎ