cap_prompt attempts

app.py

@@ -5,9 +5,9 @@ import time
 import json
 sys.path.append("./virtex/")
 
-def gen_show_caption(sub_prompt=None, cap_prompt = None):
+def gen_show_caption(sub_prompt=None, cap_prompt = ""):
     with st.spinner("Generating Caption"):
-        subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_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")
     
@@ -44,12 +44,9 @@ sub = st.sidebar.selectbox(
 st.sidebar.title("Write a Custom Prompt")
 cap_prompt = st.sidebar.text_input(
     "Leave this blank for an unbiased caption", 
-    value=None
+    value=""
 )
 
-if cap_prompt is None:
-    st.write("HAHA")
-
 if st.sidebar.button("Random Sample Image"):
     random_image = get_rand_img(sample_images)
     sample_image = None

model.py

@@ -62,6 +62,9 @@ class VirTexModel():
                                )
             subreddit_tokens = torch.tensor(subreddit_tokens, device=self.device).long()
             
+        if prompt is not "":
+            cap_tokens = self.tokenizer.encode(subreddit_tokens) 
+            subreddit_tokens = torch.cat((subreddit_tokens, cap_tokens))
             
         predictions: List[Dict[str, Any]] = []