Update app.py

app.py

@@ -2,6 +2,14 @@ import streamlit as st
 from PIL import Image
 from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel,RobertaTokenizerFast
 
+
+import requests
+from PIL import Image
+
+from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel
+
+
+
 #import torch
 #from transformers import BlipProcessor, BlipForConditionalGeneration
 
@@ -12,22 +20,22 @@ def set_page_config():
         layout='wide',
     )
 
-#def initialize_model():
-#    hf_model = "Salesforce/blip-image-captioning-large"
-#    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-#    processor = BlipProcessor.from_pretrained(hf_model)
-#    model = BlipForConditionalGeneration.from_pretrained(hf_model).to(device) # type: ignore
- #   return processor, model, device
+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")
+   return image_processor, model,tokenizer, device
 
 def upload_image():
     return st.sidebar.file_uploader("Upload an image (we aren't storing anything)", type=["jpg", "jpeg", "png"])
 
 def resize_image(image, max_width):
-    width, height = image.size
-    if width > max_width:
-        ratio = max_width / width
-        height = int(height * ratio)
-        image = image.resize((max_width, height))
+    
+    image = image.resize((max_width, height))
+    if image.mode == "L":
+        image = image.convert("RGB")
     return image
 
 def generate_caption(processor, model, device, image):