Update app.py

app.py

@@ -56,9 +56,9 @@ def ui():
         pil_image = PIL.Image.fromarray(image)
         image = preprocess(pil_image).unsqueeze(0).to(device)
 
-        option = st.selectbox('Please select the Model',('PreTrained Model','Trained Model'))
+        option = st.selectbox('Please select the Model',('Clip Captioning','Attention Decoder','VIT+GPT2'))
 
-        if option=='PreTrained Model':
+        if option=='Clip Captioning':
             with torch.no_grad():
                 prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)
                 prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
@@ -66,12 +66,12 @@ def ui():
 
             st.image(uploaded_file, width = 500, channels = 'RGB')
             st.markdown("**PREDICTION:** " + generated_text_prefix)
-        elif option=='Trained Model': 
+        elif option=='Attention Decoder': 
             out = inference(uploaded_file)
             st.image(uploaded_file, width = 500, channels = 'RGB')
             st.markdown("**PREDICTION:** " + out)
 
-        elif option=='Fine Tuned Model': 
+        elif option=='VIT+GPT2': 
             out=show_n_generate(uploaded_file, greedy = False, model = model_trained)
             st.image(uploaded_file, width = 500, channels = 'RGB')
             st.markdown("**PREDICTION:** " + out)