Add models to state

apps/mlm.py

@@ -27,7 +27,7 @@ def app(state):
 
     # @st.cache(persist=False) # TODO: Make this work with mlm_state. Currently not supported.
     def predict(transformed_image, caption_inputs):
-        outputs = model(pixel_values=transformed_image, **caption_inputs)
+        outputs = mlm_state.model(pixel_values=transformed_image, **caption_inputs)
         indices = np.where(caption_inputs["input_ids"] == bert_tokenizer.mask_token_id)[
             1
         ][0]
@@ -56,10 +56,10 @@ def app(state):
         image = plt.imread(image_path)
         mlm_state.mlm_image = image
 
-    #if model is None:
-    # Display Top-5 Predictions
-    with st.spinner("Loading model..."):
-        model = load_model(mlm_checkpoints[0])
+    if mlm_state.model is None:
+        # Display Top-5 Predictions
+        with st.spinner("Loading model..."):
+            mlm_state.model = load_model(mlm_checkpoints[0])
 
     if st.button(
         "Get a random example",

apps/vqa.py

@@ -31,7 +31,7 @@ def app(state):
     # @st.cache(persist=False)
     def predict(transformed_image, question_inputs):
         return np.array(
-            model(pixel_values=transformed_image, **question_inputs)[0][0]
+            vqa_state.model(pixel_values=transformed_image, **question_inputs)[0][0]
         )
 
     # @st.cache(persist=False)
@@ -65,11 +65,12 @@ def app(state):
         image = plt.imread(image_path)
         vqa_state.vqa_image = image
 
-    # if model is None:
+    if vqa_state.model is None:
+        with st.spinner("Loading model..."):
+            vqa_state.model = load_model(vqa_checkpoints[0])
     
     # Display Top-5 Predictions
-    with st.spinner("Loading model..."):
-        model = load_model(vqa_checkpoints[0])
+    
 
     if st.button(
         "Get a random example",