Update vlm.py

vlm.py

@@ -5,7 +5,8 @@ Author: Didier Guillevic
 Date: 2025-05-08
 """
 
-from transformers import AutoProcessor, AutoModelForImageTextToText
+from transformers import AutoProcessor
+from transformers import Mistral3ForConditionalGeneration
 import torch
 
 #
@@ -14,11 +15,11 @@ import torch
 model_id = "OPEA/Mistral-Small-3.1-24B-Instruct-2503-int4-AutoRound-awq-sym"
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForImageTextToText.from_pretrained(
+model = Mistral3ForConditionalGeneration.from_pretrained(
     model_id, 
     _attn_implementation="flash_attention_2",
     torch_dtype=torch.bfloat16
-).to(device)
+).eval().to(device)
 
 #
 # Encode images as base64