prompt fixing

app.py

@@ -13,10 +13,7 @@ from model import *
 
 def gen_show_caption(sub_prompt=None, cap_prompt = ""):
     with st.spinner("Generating Caption"):
-        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)
+        subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt, prompt=cap_prompt)
         st.markdown(
             f"""
             <style>
@@ -31,7 +28,7 @@ def gen_show_caption(sub_prompt=None, cap_prompt = ""):
                 }}
             </style>
 
-            ### <red> r/{subreddit} </red> <blue> {prompt} </blue> {caption}
+            ### <red> r/{subreddit} </red> <blue> {cap_prompt} </blue> {caption}
             """, 
             unsafe_allow_html=True)
     

model.py

@@ -117,16 +117,14 @@ class VirTexModel():
                 subreddit = "".join(subreddit.split())
                 rest_of_caption = rest_of_caption.strip()
             else:
-                subreddit, rest_of_caption = "", caption
+                subreddit, rest_of_caption = "", caption.strip()
 
             # split prompt for coloring:
             if prompt is not "":
-                _, caption = caption.split(prompt)
+                _, rest_of_caption = caption.split(prompt.strip())
             
             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():