[OCR API] Reverted.

main.py

@@ -2,76 +2,47 @@ 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"
-
-# Check for Metal GPU support on macOS
-device = "mps" if torch.backends.mps.is_available() else "cpu"
-
+min_pixels = 256*28*28
+max_pixels = 1280*28*28
 processor = AutoProcessor.from_pretrained(
     checkpoint,
-    min_pixels=256 * 28 * 28,
-    max_pixels=1280 * 28 * 28
+    min_pixels=min_pixels,
+    max_pixels=max_pixels
 )
-
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     checkpoint,
-    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
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    # attn_implementation="flash_attention_2",
 )
 
-
-# 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)
-
-    # Process image
-    image_inputs = [load_and_resize_image(image_url)]
-    video_inputs = None
-
-    # Process inputs
+    image_inputs, video_inputs = process_vision_info(messages)
     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(device)
-
-    # Generate response
+    ).to(model.device)
     with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=64)  # Reduced for faster inference
-
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
     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]}