Update app.py

app.py

@@ -12,37 +12,42 @@ Original file is located at
 from PIL import Image
 from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, PreTrainedTokenizerFast
 import requests
+from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
+import torch
+from PIL import Image
 
-model = VisionEncoderDecoderModel.from_pretrained("sachin/vit2distilgpt2")
-
-vit_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
-
-tokenizer = PreTrainedTokenizerFast.from_pretrained("distilgpt2")
+model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
-# url = 'https://d2gp644kobdlm6.cloudfront.net/wp-content/uploads/2016/06/bigstock-Shocked-and-surprised-boy-on-t-113798588-300x212.jpg'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# with Image.open(requests.get(url, stream=True).raw) as img:
-#     pixel_values = vit_feature_extractor(images=img, return_tensors="pt").pixel_values
 
-#encoder_outputs = model.generate(pixel_values.to('cpu'),num_beams=5)
 
-#generated_sentences = tokenizer.batch_decode(encoder_outputs, skip_special_tokens=True)
+max_length = 16
+num_beams = 4
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+def predict_step(image_paths):
+  images = []
+  for image_path in image_paths:
+    i_image = Image.open(image_path)
+    if i_image.mode != "RGB":
+      i_image = i_image.convert(mode="RGB")
 
-#generated_sentences
+    images.append(i_image)
 
-#naive text processing 
-#generated_sentences[0].split('.')[0]
+  pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
+  pixel_values = pixel_values.to(device)
 
-# inference function
+  output_ids = model.generate(pixel_values, **gen_kwargs)
 
-def vit2distilgpt2(img):
-  pixel_values = vit_feature_extractor(images=img, return_tensors="pt").pixel_values
-  encoder_outputs = generated_ids = model.generate(pixel_values.to('cpu'),num_beams=5)
-  generated_sentences = tokenizer.batch_decode(encoder_outputs, skip_special_tokens=True)
+  preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+  preds = [pred.strip() for pred in preds]
+  return preds
 
-  return(generated_sentences[0].split('.')[0])
+#predict_step(['/content/drive/MyDrive/caption generator/horses.png'])
 
-#!wget https://media.glamour.com/photos/5f171c4fd35176eaedb36823/master/w_2560%2Cc_limit/bike.jpg
 
 import gradio as gr
 
@@ -65,7 +70,7 @@ examples = [
 ]
 
 gr.Interface(
-    vit2distilgpt2,
+    predict_step,
     inputs,
     outputs,
     title=title,