Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoImageProcessor, AutoModelForCausalLM, BlipForConditionalGeneration, VisionEncoderDecoderModel
+from transformers import AutoProcessor, AutoTokenizer, AutoImageProcessor, AutoModelForCausalLM, BlipForConditionalGeneration, VisionEncoderDecoderModel
 import torch
 
 torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
@@ -13,6 +13,7 @@ blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image
 
 vitgpt_processor = AutoImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 vitgpt_model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+vitgpt_tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -20,12 +21,15 @@ git_model.to(device)
 blip_model.to(device)
 vitgpt_model.to(device)
 
-def generate_caption(processor, model, image):
+def generate_caption(processor, model, image, tokenizer=None):
     inputs = processor(images=image, return_tensors="pt").to(device)
     
     generated_ids = model.generate(pixel_values=inputs.pixel_values, max_length=50)
-     
-    generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+    if tokenizer is not None:
+        generated_ids = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    else:
+        generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
     return generated_caption
 
@@ -35,7 +39,7 @@ def generate_captions(image):
 
     caption_blip = generate_caption(blip_processor, blip_model, image)
 
-    caption_vitgpt = generate_caption(vitgpt_processor, vitgpt_model, image)
+    caption_vitgpt = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer)
 
     return caption_git, caption_blip, caption_vitgpt