Add batching using vllm

main.py

@@ -77,6 +77,9 @@ import logging
 from typing import List
 import time
 import numpy as np
+from vllm import LLM, SamplingParams
+import torch._dynamo
+torch._dynamo.config.suppress_errors = True
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -97,6 +100,7 @@ app = FastAPI()
 # Global variables for model and processor
 model = None
 processor = None
+llm = None
 
 def load_model():
     """Load model and processor when needed"""
@@ -114,6 +118,15 @@ def load_model():
         processor = PaliGemmaProcessor.from_pretrained(model_id)
         logger.info("Model loaded successfully")
 
+def load_vllm_model():
+    global llm
+    if llm is None:
+        llm = LLM(
+            model="google/paligemma2-3b-mix-448",
+            trust_remote_code=True,
+            max_model_len=4096,
+            dtype="float16",
+        )
 def clean_memory():
     """Force garbage collection and clear CUDA cache"""
     gc.collect()
@@ -165,6 +178,37 @@ async def extract_text(background_tasks: BackgroundTasks, file: UploadFile = Fil
         logger.error(f"Error processing image: {str(e)}")
         return {"error": str(e)}
 
+@app.post("/batch_extract_text_vllm")
+async def batch_extract_text_vllm(background_tasks: BackgroundTasks, files: List[UploadFile] = File(...)):
+    try:
+        start_time = time.time()
+        load_vllm_model()
+        results = []        
+        sampling_params = SamplingParams(temperature=0.0,max_tokens=32)
+        # Load images
+        images = []
+        for file in files:
+            image_data = await file.read()
+            img = Image.open(BytesIO(image_data)).convert("RGB")
+            images.append(img)
+        for image in images:    
+            inputs = {
+                "prompt": "ocr",
+                "multi_modal_data": {
+                    "image": image
+                },
+            }
+            outputs = llm.generate(inputs, sampling_params)
+            for o in outputs:
+                generated_text = o.outputs[0].text        
+                results.append(generated_text)
+                        
+        logger.info(f"vLLM Batch processing completed in {time.time() - start_time:.2f} seconds")
+        return {"extracted_texts": results}
+    except Exception as e:
+        logger.error(f"Error in batch processing vLLM: {str(e)}")
+        return {"error": str(e)}
+
 @app.post("/batch_extract_text")
 async def batch_extract_text(batch_size:int, background_tasks: BackgroundTasks, files: List[UploadFile] = File(...)):
     """Extract text from multiple images with batching"""

requirements.txt

@@ -7,4 +7,5 @@ transformers
 torch
 accelerate
 pillow
-python-multipart
+python-multipart
+vllm