updated prompt splitting

app.py

@@ -13,21 +13,25 @@ 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 we default to /r/pics")
-        subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt, prompt = cap_prompt)
+        if cap_prompt is not "":
+            subreddit, prompt, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt, prompt = cap_prompt)
+        else:
+            subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt, prompt = cap_prompt)
         st.markdown(
             f"""
             <style>
                 red{{
                     color:#c62828
                 }}
+                blue{{
+                    color:#2a72d5
+                }}
                 mono{{
                     font-family: "Inconsolata";
                 }}
             </style>
 
-            ### <red> r/{subreddit} </red>  {caption}
+            ### <red> r/{subreddit} </red> <blue> {prompt} </blue> {caption}
             """, 
             unsafe_allow_html=True)
     
@@ -109,31 +113,30 @@ if advanced:
     nuc_size = st.sidebar.slider("Nucelus Size:", min_value=0.0, max_value=1.0, value=0.8, step=0.05)
     virtexModel.model.decoder.nucleus_size = nuc_size
 
-if False: #uploaded_image is None:# and submitted:
-    st.write("Please select a file to upload")
+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_file = sample_image
+    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")
 
-    for i in range(num_captions):
-        gen_show_caption(sub, imageLoader.text_transform(cap_prompt))
+if sub is None and imageLoader.text_transform(cap_prompt) is not "":
+    st.write("Without a specified subreddit we default to /r/pics")
+for i in range(num_captions):
+    gen_show_caption(sub, imageLoader.text_transform(cap_prompt))

model.py

@@ -92,7 +92,6 @@ class VirTexModel():
             subreddit_tokens = torch.cat(
                 [
                     subreddit_tokens,
-                    torch.tensor([self.tokenizer.token_to_id("[SEP]")], device=self.device).long(),
                     cap_tokens
                 ])
             
@@ -119,10 +118,15 @@ class VirTexModel():
                 rest_of_caption = rest_of_caption.strip()
             else:
                 subreddit, rest_of_caption = "", caption
+
+            # split prompt for coloring:
+            if prompt is not "":
+                _, caption = caption.split(prompt)
             
             is_valid_subreddit = subreddit in self.valid_subs
             
-            
+        if prompt is not "":
+            return subreddit, prompt, rest_of_caption
         return subreddit, rest_of_caption
 
 def download_files():