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',('Model', 'COCO Model','PreTrained Model','Fine Tuned Model'))
+        option = st.selectbox('Please select the Model',(' PreTrained Model','Trained Model','Fine Tuned Model'))
 
-        if option=='Model':
+        if option=='PreTrained Model':
             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,17 +66,7 @@ def ui():
 
             st.image(uploaded_file, width = 500, channels = 'RGB')
             st.markdown("**PREDICTION:** " + generated_text_prefix)
-        
-        elif option=='COCO Model':
-            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)
-            generated_text_prefix = generate2(coco_model, tokenizer, embed=prefix_embed)
-
-            st.image(uploaded_file, width = 500, channels = 'RGB')
-            st.markdown("**PREDICTION:** " + generated_text_prefix)
-
-        elif option=='PreTrained Model': 
+        elif option=='Trained Model': 
             out = inference(uploaded_file)
             st.image(uploaded_file, width = 500, channels = 'RGB')
             st.markdown("**PREDICTION:** " + out)