Update app.py

app.py

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