added caching to loading

app.py

@@ -3,6 +3,7 @@ import io
 import sys
 import time
 import json
+from model import *
 sys.path.append("./virtex/")
 
 def gen_show_caption(sub_prompt=None, cap_prompt = ""):
@@ -12,7 +13,13 @@ def gen_show_caption(sub_prompt=None, cap_prompt = ""):
         subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt, prompt = cap_prompt)
         st.header("Predicted Caption:\n\n")
 #         st.subheader(f"r/{subreddit}:\t{caption}\n")
-        st.markdown(f"## r/ {subreddit} \t **{caption}**")
+        st.markdown(
+            f"""
+            r/{subreddit}
+            \t
+            **{caption}**
+            """
+        )
     
 
 st.title("Image Captioning Demo from Redcaps")
@@ -23,12 +30,7 @@ st.sidebar.markdown(
 )
 
 with st.spinner("Loading Model"):
-    from model import *
-    sample_images = get_samples()
-    virtexModel = VirTexModel()
-    imageLoader = ImageLoader()
-    valid_subs = json.load(open(VALID_SUBREDDITS_PATH))
-    valid_subs.insert(0, None)
+    virtexModel, imageLoader, sample_images, valid_subs = create_objects()
 
 random_image = get_rand_img(sample_images)
 rand_idx = 0

model.py

@@ -1,3 +1,4 @@
+import streamlit as st
 from huggingface_hub import hf_hub_url, cached_download
 from PIL import Image
 import os
@@ -127,3 +128,12 @@ def get_rand_img(samples):
     i = random.randint(0,len(samples)-1)
     return i, samples[i]
 
+@st.cache
+def create_objects():
+    sample_images = get_samples()
+    virtexModel = VirTexModel()
+    imageLoader = ImageLoader()
+    valid_subs = json.load(open(VALID_SUBREDDITS_PATH))
+    valid_subs.insert(0, None)
+    return virtexModel, imageLoader, sample_images, valid_subs
+