code in comments pre-staging

app.py

@@ -5,6 +5,7 @@ import requests
 import matplotlib.pyplot as plt
 import gradio as gr
 from transformers import ImageClassificationPipeline, PerceiverForImageClassificationConvProcessing, PerceiverFeatureExtractor
+from transformers import VisionEncoderDecoderModel
 from transformers import AutoTokenizer
 import torch
 
@@ -27,12 +28,25 @@ model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/
 image_pipe = ImageClassificationPipeline(model=model, feature_extractor=feature_extractor)
 
 '''
-repo_name = "ydshieh/vit-gpt2-coco-en"
 
+
+# initialize a vit-bert from a pretrained ViT and a pretrained BERT model. Note that the cross-attention layers will be randomly initialized
+model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained(
+    "google/vit-base-patch16-224-in21k", "bert-base-uncased"
+)
+# saving model after fine-tuning
+model.save_pretrained("./vit-bert")
+# load fine-tuned model
+model = VisionEncoderDecoderModel.from_pretrained("./vit-bert")
+
+
+
+repo_name = "ydshieh/vit-gpt2-coco-en"
+test_image = "cats.jpg"
 feature_extractor2 = ViTFeatureExtractor.from_pretrained(repo_name)
 tokenizer = AutoTokenizer.from_pretrained(repo_name)
 model2 = VisionEncoderDecoderModel.from_pretrained(repo_name)
-pixel_values = feature_extractor2(image, return_tensors="pt").pixel_values
+pixel_values = feature_extractor2(test_image, return_tensors="pt").pixel_values
 
 # autoregressively generate text (using beam search or other decoding strategy)
 generated_ids = model2.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True)