Update app.py

app.py

@@ -1,5 +1,16 @@
 # !pip install torch
 # import torch
+from transformers import AutoTokenizer, AutoModel ,AutoConfig
+import torch
+from transformers import  ViTImageProcessor, VisionEncoderDecoderModel,RobertaTokenizerFast
+import PIL 
+
+
+# Move model to GPU , depnding on device
+device2 = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Load the model
+
+
 
 import streamlit as st
 from PIL import Image
@@ -24,10 +35,10 @@ def set_page_config():
 
 def initialize_model():
     device = 'cpu'
-    # load a fine-tuned image captioning model and corresponding tokenizer and image processor
-    model = AutoModel.from_pretrained("sourabhbargi11/Caption_generator_model").to(device) 
-    tokenizer = AutoTokenizer.from_pretrained("sourabhbargi11/Caption_generator_model")
-    image_processor = ViTImageProcessor.from_pretrained("sourabhbargi11/Caption_generator_model")
+    config = AutoConfig.from_pretrained("sourabhbargi11/Caption_generator_model")  
+    model = VisionEncoderDecoderModel.from_pretrained("sourabhbargi11/Caption_generator_model", config=config).to(device)
+    tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
+    image_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224",device=device)
     return image_processor, model,tokenizer, device
 
 def upload_image():