Fix app.py: lazy loading, correct Qwen2.5-VL class, qwen_vl_utils

app.py

@@ -1,25 +1,96 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForSeq2SeqLM
 import torch
+import os
 from PIL import Image
 
-# Load your model
-processor = AutoProcessor.from_pretrained("OdiaGenAIOCR/odia-ocr-qwen-finetuned_v3")
-model = AutoModelForSeq2SeqLM.from_pretrained("OdiaGenAIOCR/odia-ocr-qwen-finetuned_v3")
+MODEL_ID = "OdiaGenAIOCR/odia-ocr-qwen-finetuned_v3"
+CHECKPOINT = "checkpoint-1300"
+
+_model = None
+_processor = None
+
+def load_model():
+    global _model, _processor
+    if _model is not None:
+        return _model, _processor
+
+    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+    from peft import PeftModel
+
+    BASE_MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    _processor = AutoProcessor.from_pretrained(
+        MODEL_ID,
+        revision=CHECKPOINT,
+        trust_remote_code=True,
+    )
+
+    _model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        MODEL_ID,
+        revision=CHECKPOINT,
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        device_map="auto" if device == "cuda" else None,
+        trust_remote_code=True,
+    )
+    if device == "cpu":
+        _model = _model.to(device)
+    _model.eval()
+    return _model, _processor
+
 
 def ocr_predict(image: Image.Image):
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-    return text
+    if image is None:
+        return "Please upload an image."
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    try:
+        model, processor = load_model()
+    except Exception as e:
+        return f"Model load error: {e}"
+
+    from qwen_vl_utils import process_vision_info
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image},
+                {"type": "text", "text": "Please perform OCR on this image and output only the recognized text."},
+            ],
+        }
+    ]
+
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    ).to(device)
+
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs, max_new_tokens=512)
+
+    generated_ids_trimmed = [
+        out_ids[len(in_ids):]
+        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return output_text[0]
+
 
-# Gradio interface
 iface = gr.Interface(
     fn=ocr_predict,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(),
+    outputs=gr.Textbox(label="Recognized Odia Text"),
     title="OdiaGenAI OCR",
-    description="Upload Odia images or scanned documents for OCR."
+    description="Upload Odia images or scanned documents for OCR using the fine-tuned Qwen2.5-VL model.",
+    examples=[],
 )
 
 iface.launch()