[OCR API] Optimization pass.

main.py

@@ -2,47 +2,76 @@ from fastapi import FastAPI, Query
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
 import torch
+import requests
+from PIL import Image
+from io import BytesIO
 
 app = FastAPI()
 
+# Load model and processor
 checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
-min_pixels = 256*28*28
-max_pixels = 1280*28*28
+
+# Check for Metal GPU support on macOS
+device = "mps" if torch.backends.mps.is_available() else "cpu"
+
 processor = AutoProcessor.from_pretrained(
     checkpoint,
-    min_pixels=min_pixels,
-    max_pixels=max_pixels
+    min_pixels=256 * 28 * 28,
+    max_pixels=1280 * 28 * 28
 )
+
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     checkpoint,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    # attn_implementation="flash_attention_2",
+    torch_dtype=torch.float16 if device == "mps" else torch.bfloat16,  # Use float16 on Apple Metal
+    device_map={"": 0} if device == "mps" else "cpu",
+    attn_implementation="flash_attention_2",  # If it supports Mac
 )
 
+
+# Function to load and resize images (reduces processing time)
+def load_and_resize_image(image_url):
+    response = requests.get(image_url)
+    image = Image.open(BytesIO(response.content)).convert("RGB")
+    image = image.resize((512, 512))  # Resize to 512x512 to speed up processing
+    return image
+
+
 @app.get("/")
 def read_root():
     return {"message": "API is live. Use the /predict endpoint."}
 
+
 @app.get("/predict")
 def predict(image_url: str = Query(...), prompt: str = Query(...)):
     messages = [
         {"role": "system", "content": "You are a helpful assistant with vision abilities."},
         {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
     ]
+
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
+
+    # Process image
+    image_inputs = [load_and_resize_image(image_url)]
+    video_inputs = None
+
+    # Process inputs
     inputs = processor(
         text=[text],
         images=image_inputs,
         videos=video_inputs,
         padding=True,
+        truncation=True,  # Ensures token limit
+        max_length=512,  # Prevents excessive memory usage
         return_tensors="pt",
-    ).to(model.device)
+    ).to(device)
+
+    # Generate response
     with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=128)
+        generated_ids = model.generate(**inputs, max_new_tokens=64)  # Reduced for faster inference
+
     generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
     output_texts = processor.batch_decode(
         generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )
+
     return {"response": output_texts[0]}