Revert "Remove session handling hack."

This reverts commit 1985070cad8718dcbf1006e277626c52d3a12c86.

app/app.py

@@ -1,11 +1,31 @@
 #!/usr/bin/env python
 # coding: utf-8
 
+import random
 from dalle_mini.backend import ServiceError, get_images_from_backend
+
 import streamlit as st
 
+# streamlit.session_state is not available in Huggingface spaces.
+# Session state hack https://huggingface.slack.com/archives/C025LJDP962/p1626527367443200?thread_ts=1626525999.440500&cid=C025LJDP962
+
+from streamlit.report_thread import get_report_ctx
+def query_cache(q_emb=None):
+    ctx = get_report_ctx()
+    session_id = ctx.session_id
+    session = st.server.server.Server.get_current()._get_session_info(session_id).session
+    if not hasattr(session, "_query_state"):
+        setattr(session, "_query_state", q_emb)
+    if q_emb:
+        session._query_state = q_emb
+    return session._query_state
+
+def set_run_again(state):
+    query_cache(state)
+
 def should_run_again():
-    return st.session_state.get("again", False)
+    state = query_cache()
+    return state if state is not None else False
 
 st.sidebar.markdown("""
 <style>
@@ -37,7 +57,7 @@ prompt = st.text_input("What do you want to see?")
 
 test = st.empty()
 DEBUG = False
-if prompt != "" or (should_run_again() and prompt != ""):
+if prompt != "" or (should_run_again and prompt != ""):
     container = st.empty()
     # The following mimics `streamlit.info()`.
     # I tried to get the secondary background color using `components.streamlit.config.get_options_for_section("theme")["secondaryBackgroundColor"]`
@@ -71,7 +91,8 @@ if prompt != "" or (should_run_again() and prompt != ""):
             cols[i%4].image(img)
 
         container.markdown(f"**{prompt}**")
-        st.session_state["again"] = st.button('Again!', key='again_button')
+        
+        set_run_again(st.button('Again!', key='again_button'))
     
     except ServiceError as error:
         container.text(f"Service unavailable, status: {error.status_code}")