Update app.py

app.py

@@ -1,13 +1,15 @@
 import torch 
 import re 
 import gradio as gr
-from transformers import AutoTokenizer, ViTFeatureExtractor, AutoModel
+from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel
 
 device='cpu'
 encoder_checkpoint = "Thibalte/captionning_project"
+decoder_checkpoint = "Thibalte/captionning_project"
+model_checkpoint = "Thibalte/captionning_project"
 feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
-tokenizer = AutoTokenizer.from_pretrained("Thibalte/captionning_project")
-model = AutoModel.from_pretrained("Thibalte/captionning_project")
+tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
+model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
 
 
 def predict(image,max_length=24, num_beams=4):