Update app.py

app.py

@@ -15,13 +15,14 @@ from PIL import Image, ImageOps
 import requests
 
 from transformers import (
-    AutoTokenizer,
+    Qwen2VLForConditionalGeneration,
+    Qwen2_5_VLForConditionalGeneration,
+    AutoModelForImageTextToText,
     AutoProcessor,
     TextIteratorStreamer,
 )
+
 from transformers.image_utils import load_image
-# The custom model class is imported via trust_remote_code=True
-from transformers import AutoModelForImageTextToText
 
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
@@ -116,12 +117,10 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load Nanonets-OCR2-3B
 MODEL_ID_3B = "nanonets/Nanonets-OCR2-3B"
 processor_3b = AutoProcessor.from_pretrained(MODEL_ID_3B, trust_remote_code=True)
-model_3b = AutoModelForImageTextToText.from_pretrained(
+model_3b = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID_3B,
     dtype=torch.float16,
-    #device_map="auto",
     trust_remote_code=True,
-    attn_implementation="flash_attention_2"
 ).to(device).eval()
 
 # Load Nanonets-OCR2-1.5B-exp
@@ -130,7 +129,6 @@ processor_1_5b = AutoProcessor.from_pretrained(MODEL_ID_1_5B, trust_remote_code=
 model_1_5b = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID_1_5B,
     dtype=torch.float16,
-    #device_map="auto",
     trust_remote_code=True,
     attn_implementation="flash_attention_2"
 ).to(device).eval()