default to pics for subreddit sampling

app.py

@@ -9,7 +9,7 @@ from model import *
 def gen_show_caption(sub_prompt=None, cap_prompt = ""):
     with st.spinner("Generating Caption"):
         if sub_prompt is None and cap_prompt is not "":
-            st.write("Without a specified subreddit, caption prompts will skip subreddit prediction")
+            st.write("Without a specified subreddit we default to /r/pics")
         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")
@@ -31,7 +31,7 @@ st.sidebar.markdown(
     You can also generate captions as if they are from specific subreddits,
     as if they start with a particular prompt, or even both.
     
-    Feel free to share your results on twitter with #redcaps or with a friend.
+    Share your results on twitter with #redcaps or with a friend.
     """
 )
 
@@ -39,91 +39,91 @@ with st.spinner("Loading Model"):
     virtexModel, imageLoader, sample_images, valid_subs = create_objects()
     
 
-staggered = st.sidebar.checkbox("Iteratively Generate Captions")
+# staggered = st.sidebar.checkbox("Iteratively Generate Captions")
 
-if staggered:
-    pass
-else:
-    
-    select_idx = None
+# if staggered:
+#     pass
+# else:
+
+select_idx = None
 
-    st.sidebar.title("Select a sample image")
+st.sidebar.title("Select a sample image")
 
-    if st.sidebar.button("Random Sample Image"):
-        select_idx = get_rand_idx(sample_images)
+if st.sidebar.button("Random Sample Image"):
+    select_idx = get_rand_idx(sample_images)
 
-    sample_image = sample_images[0 if select_idx is None else select_idx]
+sample_image = sample_images[0 if select_idx is None else select_idx]
 
 
-    uploaded_image = None
-    with st.sidebar.form("file-uploader-form", clear_on_submit=True):
-        uploaded_file = st.file_uploader("Choose a file")
-        submitted = st.form_submit_button("Submit")
-        if uploaded_file is not None and submitted:
-            uploaded_image = Image.open(io.BytesIO(uploaded_file.getvalue()))
-            select_idx = None # set this to help rewrite the cache
+uploaded_image = None
+with st.sidebar.form("file-uploader-form", clear_on_submit=True):
+    uploaded_file = st.file_uploader("Choose a file")
+    submitted = st.form_submit_button("Submit")
+    if uploaded_file is not None and submitted:
+        uploaded_image = Image.open(io.BytesIO(uploaded_file.getvalue()))
+        select_idx = None # set this to help rewrite the cache
 
-    # class OnChange():
-    #     def __init__(self, idx):
-    #         self.idx = idx
+# class OnChange():
+#     def __init__(self, idx):
+#         self.idx = idx
 
-    #     def __call__(self):
-    #         st.write(f"the idx is: {self.idx}")
-    #         st.write(f"the sample_image is {sample_image}")
+#     def __call__(self):
+#         st.write(f"the idx is: {self.idx}")
+#         st.write(f"the sample_image is {sample_image}")
 
-    # sample_image = st.sidebar.selectbox(
-    #     "",
-    #     sample_images,
-    #     index = 0 if select_idx is None else select_idx,
-    #     on_change=OnChange(0 if select_idx is None else select_idx)
-    # )
+# sample_image = st.sidebar.selectbox(
+#     "",
+#     sample_images,
+#     index = 0 if select_idx is None else select_idx,
+#     on_change=OnChange(0 if select_idx is None else select_idx)
+# )
 
-    st.sidebar.title("Select a Subreddit")
-    sub = st.sidebar.selectbox(
-        "Select None for a Predicted Subreddit",
-        valid_subs
-    )
+st.sidebar.title("Select a Subreddit")
+sub = st.sidebar.selectbox(
+    "Type below to condition on a subreddit. Select None for a predicted subreddit",
+    valid_subs
+)
 
-    st.sidebar.title("Write a Custom Prompt")
-    cap_prompt = st.sidebar.text_input(
-        "Leave this blank for an unbiased caption", 
-        value=""
-    )
+st.sidebar.title("Write a Custom Prompt")
+cap_prompt = st.sidebar.text_input(
+    "Write the start of your caption below", 
+    value=""
+)
+
+_ = st.sidebar.button("Regenerate Caption")
 
-    _ = st.sidebar.button("Regenerate Caption")
-    
 #     advanced = st.sidebar.checkbox("Advanced Options")
-    
+
 #     if advanced:
 #         nuc_size = st.sidebar.slider("")
 
-    if uploaded_image is None and submitted:
-        st.write("Please select a file to upload")
+if uploaded_image is None and submitted:
+    st.write("Please select a file to upload")
 
-    else:
-        image_file = sample_image
+else:
+    image_file = sample_image
 
-        # LOAD AND CACHE THE IMAGE
-        if uploaded_image is not None:
-            image = uploaded_image
-        elif select_idx is None and 'image' in st.session_state:
-            image = st.session_state['image']
-        else:
-            image = Image.open(image_file)
+    # LOAD AND CACHE THE IMAGE
+    if uploaded_image is not None:
+        image = uploaded_image
+    elif select_idx is None and 'image' in st.session_state:
+        image = st.session_state['image']
+    else:
+        image = Image.open(image_file)
 
-        image = image.convert("RGB")
+    image = image.convert("RGB")
 
-        st.session_state['image'] = image
+    st.session_state['image'] = image
 
 
-        image_dict = imageLoader.transform(image)
+    image_dict = imageLoader.transform(image)
 
-        show_image = imageLoader.show_resize(image)
+    show_image = imageLoader.show_resize(image)
 
-        show = st.image(show_image)
-        show.image(show_image, "Your Image")
+    show = st.image(show_image)
+    show.image(show_image, "Your Image")
 
-        gen_show_caption(sub, imageLoader.text_transform(cap_prompt))
+    gen_show_caption(sub, imageLoader.text_transform(cap_prompt))
 
     # from model import *
     # sample_images = get_samples()

model.py

@@ -80,15 +80,19 @@ class VirTexModel():
         if prompt is not "":
             # at present prompts without subreddits will break without this change
             # TODO FIX
-            if True: #sub_prompt is not None:
-                cap_tokens = self.tokenizer.encode(prompt)
-                cap_tokens = torch.tensor(cap_tokens, device=self.device).long()
-                subreddit_tokens = torch.cat(
-                    [
-                        subreddit_tokens,
-                        torch.tensor([self.tokenizer.token_to_id("[SEP]")], device=self.device).long(),
-                        cap_tokens
-                    ])
+            cap_tokens = self.tokenizer.encode(prompt)
+            cap_tokens = torch.tensor(cap_tokens, device=self.device).long()
+            subreddit_tokens = subreddit_tokens if sub_prompt is not None else torch.tensor([
+                [self.model.sos_index] +
+                self.tokenizer.encode("pics") +
+                [self.tokenizer.token_to_id("[SEP]")]
+            ])
+            subreddit_tokens = torch.cat(
+                [
+                    subreddit_tokens,
+                    torch.tensor([self.tokenizer.token_to_id("[SEP]")], device=self.device).long(),
+                    cap_tokens
+                ])
             
             
         predictions: List[Dict[str, Any]] = []