Update app.py

app.py

@@ -25,9 +25,12 @@ with st.sidebar:
     api_key = st.text_input("Cohere API Key", type="password")
     selected_model = st.selectbox("Model", options=list(model_info.keys()))
     if selected_model.startswith("c4ai-aya-vision"):
-        uploaded = st.file_uploader("Upload image", type=["png","jpg","jpeg"], key="uploaded")
+        uploaded_file = st.file_uploader("Upload image", type=["png", "jpg", "jpeg"])
+        if uploaded_file:
+            st.session_state.image_data = uploaded_file.read()
     if st.button("Clear Chat"):
         st.session_state.clear()
+        st.session_state.image_data = None
         st.rerun()
     st.divider()
     st.subheader(selected_model)
@@ -40,9 +43,11 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 if "first_message_sent" not in st.session_state:
     st.session_state.first_message_sent = False
+if "image_data" not in st.session_state:
+    st.session_state.image_data = None
 
 if not st.session_state.first_message_sent:
-    st.markdown("<h1 style='text-align: center; color: #4a4a4a; margin-top: 100px;'>How can Cohere help you today?</h1>", unsafe_allow_html=True)
+    st.markdown("<h1 style='text-align:center;color:#4a4a4a;margin-top:100px;'>How can Cohere help you today?</h1>", unsafe_allow_html=True)
 
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"], avatar=USER_PFP if msg["role"] == "user" else AI_PFP):
@@ -60,31 +65,19 @@ if prompt := st.chat_input("Message..."):
     if not api_key:
         st.error("API key required")
         st.stop()
-
     st.session_state.first_message_sent = True
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user", avatar=USER_PFP):
         st.markdown(prompt)
-
     co = cohere.ClientV2(api_key)
     content = [{"type": "text", "text": prompt}]
-    if selected_model.startswith("c4ai-aya-vision") and st.session_state.get("uploaded"):
-        data = st.session_state.uploaded.read()
-        data_url = f"data:image/jpeg;base64,{base64.b64encode(data).decode('utf-8')}"
+    if selected_model.startswith("c4ai-aya-vision") and st.session_state.image_data:
+        data_url = "data:image/jpeg;base64," + base64.b64encode(st.session_state.image_data).decode('utf-8')
         content.append({"type": "image_url", "image_url": {"url": data_url}})
-
-    response = co.chat(
-        model=selected_model,
-        messages=[*st.session_state.messages, {"role": "user", "content": content}],
-        temperature=0.3
-    )
+    response = co.chat(model=selected_model, messages=[*st.session_state.messages, {"role": "user", "content": content}], temperature=0.3)
     items = response.message.content
     reply = "".join(getattr(i, 'text', '') for i in items)
-
     with st.chat_message("assistant", avatar=AI_PFP):
         st.markdown(reply)
-
     st.session_state.messages.append({"role": "assistant", "content": items})
-
-    if "uploaded" in st.session_state:
-        st.session_state.uploaded = None
+    st.session_state.image_data = None