Update app.py

app.py

@@ -1,14 +1,19 @@
-
-!pip install torch
-
-import torch
+# !pip install torch
+# import torch
 
 import streamlit as st
 from PIL import Image
-from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel,RobertaTokenizerFast, VisionEncoderDecoderModel
+# from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel,RobertaTokenizerFast, VisionEncoderDecoderModel
+#from transformers import BlipProcessor, BlipForConditionalGeneration
+
+
+# Load model directly
+from transformers import AutoTokenizer, AutoModel
+
+# tokenizer = AutoTokenizer.from_pretrained("sourabhbargi11/Caption_generator_model")
+# model = AutoModel.from_pretrained("sourabhbargi11/Caption_generator_model")
 
 
-#from transformers import BlipProcessor, BlipForConditionalGeneration
 
 def set_page_config():
     st.set_page_config(
@@ -20,9 +25,9 @@ def set_page_config():
 def initialize_model():
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     # load a fine-tuned image captioning model and corresponding tokenizer and image processor
-    model = VisionEncoderDecoderModel.from_pretrained("sourabhbargi11/caption4").to(device) 
-    tokenizer = GPT2TokenizerFast.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
-    image_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    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")
     return image_processor, model,tokenizer, device
 
 def upload_image():