Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from PIL import Image
 import clipGPT
-# import vitGPT
+import vitGPT
 import skimage.io as io
 import PIL.Image
 
@@ -11,9 +11,9 @@ def generate_caption_clipgpt(image):
     caption = clipGPT.generate_caption_clipgpt(image)
     return caption
 
-# def generate_caption_vitgpt(image):
-#     caption = vitGPT.generate_caption(image)
-#     return caption
+def generate_caption_vitgpt(image):
+    caption = vitGPT.generate_caption(image)
+    return caption
 
 
 
@@ -48,6 +48,8 @@ with gr.Blocks() as demo:
     def predict(img, model_name):
         if model_name == "CLIP-GPT2":
             return generate_caption_clipgpt(img)  
+        elif model_name == "ViT-GPT2"
+            return generate_caption_vitgpt(img)
         else:
             return "Caption generation for this model is not yet implemented."