Fix image display issues and add caching

app.py

@@ -46,7 +46,7 @@ code_to_name = {
 @st.cache
 def generate_sequence(pixel_values, lang_code, num_beams, temperature, top_p):
     lang_code = language_mapping[lang_code]
-    output_ids = model.generate(input_ids=pixel_values, forced_bos_token_id=tokenizer.lang_code_to_id[lang_code], max_length=64, num_beams=num_beams, temperature=temperature, top_p = top_p)
+    output_ids = state.model.generate(input_ids=pixel_values, forced_bos_token_id=tokenizer.lang_code_to_id[lang_code], max_length=64, num_beams=num_beams, temperature=temperature, top_p = top_p)
     print(output_ids)
     output_sequence = tokenizer.batch_decode(output_ids[0], skip_special_tokens=True, max_length=64)
     return output_sequence
@@ -75,7 +75,8 @@ st.sidebar.title("Generation Parameters")
 num_beams = st.sidebar.number_input(label="Number of Beams", min_value=2, max_value=10, value=4, step=1, help="Number of beams to be used in beam search.")
 temperature = st.sidebar.select_slider(label="Temperature", options = list(np.arange(0.0,1.1, step=0.1)), value=1.0, help ="The value used to module the next token probabilities.", format_func=lambda x: f"{x:.2f}")
 top_p = st.sidebar.select_slider(label = "Top-P", options = list(np.arange(0.0,1.1, step=0.1)),value=1.0, help="Nucleus Sampling : If set to float < 1, only the most probable tokens with probabilities that add up to :obj:`top_p` or higher are kept for generation.", format_func=lambda x: f"{x:.2f}")
-
+if st.sidebar.button("Clear All Cache"):
+    caching.clear_cache()
 image_col, intro_col = st.beta_columns([3, 8])
 image_col.image("./misc/mic-logo.png", use_column_width="always")
 intro_col.write(read_markdown("intro.md"))
@@ -98,8 +99,11 @@ with st.beta_expander("Article"):
     # st.write(read_markdown("checkpoints.md"))
     st.write(read_markdown("acknowledgements.md"))
 
+if state.model is None:
+    with st.spinner("Loading model..."):
+        state.model = load_model(checkpoints[0])
 
-first_index = 20
+first_index = 25
 # Init Session State
 if state.image_file is None:
     state.image_file = dummy_data.loc[first_index, "image_file"]
@@ -110,9 +114,10 @@ if state.image_file is None:
     image = plt.imread(image_path)
     state.image = image
 
-col1, col2 = st.beta_columns([5, 5])
+new_col1, new_col2 = st.beta_columns([5,5])
 
-if col1.button("Get a random example", help="Get a random example from one of the seeded examples."):
+
+if new_col2.button("Get a random example", help="Get a random example from one of the seeded examples."):
     sample = dummy_data.sample(1).reset_index()
     state.image_file = sample.loc[0, "image_file"]
     state.caption = sample.loc[0, "caption"].strip("- ")
@@ -122,21 +127,10 @@ if col1.button("Get a random example", help="Get a random example from one of th
     image = plt.imread(image_path)
     state.image = image
 
-if col2.sidebar.button("Clear All Cache"):
-    caching.clear_cache()
-
-# uploaded_file = col2.file_uploader("Upload your image", type=["png", "jpg", "jpeg"])
-# if uploaded_file is not None:
-#     state.image_file = os.path.join("images", uploaded_file.name)
-#     state.image = np.array(Image.open(uploaded_file))
-
 transformed_image = get_transformed_image(state.image)
 
-new_col1, new_col2 = st.beta_columns([5,5])
 # Display Image
 new_col1.image(state.image, use_column_width="always")
-
-
 # Display Reference Caption
 new_col2.write("**Reference Caption**: " + state.caption)
 new_col2.markdown(
@@ -153,8 +147,6 @@ lang_id = new_col2.selectbox(
     help="The language in which caption is to be generated."
 )
 
-with st.spinner("Loading model..."):
-    model = load_model(checkpoints[0])
 sequence = ['']
 if new_col2.button("Generate Caption", help="Generate a caption in the specified language."):
     with st.spinner("Generating Sequence..."):