Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
 import torch
-from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers import AutoProcessor, AutoModelForImageTextToText
 from PIL import Image
 
 # Load model and processor
 model_id = "ibm-granite/granite-docling-258M"
 # We must use trust_remote_code=True for this specific model architecture
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForVision2Seq.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32)
+model = AutoModelForImageTextToText.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32)
 
 def process_image(image):
     # Convert image to RGB